Add files via upload

TTS/tts/models/__init__.py

@@ -0,0 +1,14 @@
+from typing import Dict, List, Union
+
+from TTS.utils.generic_utils import find_module
+
+
+def setup_model(config: "Coqpit", samples: Union[List[List], List[Dict]] = None) -> "BaseTTS":
+    print(" > Using model: {}".format(config.model))
+    # fetch the right model implementation.
+    if "base_model" in config and config["base_model"] is not None:
+        MyModel = find_module("TTS.tts.models", config.base_model.lower())
+    else:
+        MyModel = find_module("TTS.tts.models", config.model.lower())
+    model = MyModel.init_from_config(config=config, samples=samples)
+    return model

TTS/tts/models/align_tts.py

@@ -0,0 +1,448 @@
+from dataclasses import dataclass, field
+from typing import Dict, List, Union
+
+import torch
+from coqpit import Coqpit
+from torch import nn
+
+from TTS.tts.layers.align_tts.mdn import MDNBlock
+from TTS.tts.layers.feed_forward.decoder import Decoder
+from TTS.tts.layers.feed_forward.duration_predictor import DurationPredictor
+from TTS.tts.layers.feed_forward.encoder import Encoder
+from TTS.tts.layers.generic.pos_encoding import PositionalEncoding
+from TTS.tts.models.base_tts import BaseTTS
+from TTS.tts.utils.helpers import generate_path, maximum_path, sequence_mask
+from TTS.tts.utils.speakers import SpeakerManager
+from TTS.tts.utils.text.tokenizer import TTSTokenizer
+from TTS.tts.utils.visual import plot_alignment, plot_spectrogram
+from TTS.utils.io import load_fsspec
+
+
+@dataclass
+class AlignTTSArgs(Coqpit):
+    """
+    Args:
+        num_chars (int):
+            number of unique input to characters
+        out_channels (int):
+            number of output tensor channels. It is equal to the expected spectrogram size.
+        hidden_channels (int):
+            number of channels in all the model layers.
+        hidden_channels_ffn (int):
+            number of channels in transformer's conv layers.
+        hidden_channels_dp (int):
+            number of channels in duration predictor network.
+        num_heads (int):
+            number of attention heads in transformer networks.
+        num_transformer_layers (int):
+            number of layers in encoder and decoder transformer blocks.
+        dropout_p (int):
+            dropout rate in transformer layers.
+        length_scale (int, optional):
+            coefficient to set the speech speed. <1 slower, >1 faster. Defaults to 1.
+        num_speakers (int, optional):
+            number of speakers for multi-speaker training. Defaults to 0.
+        external_c (bool, optional):
+            enable external speaker embeddings. Defaults to False.
+        c_in_channels (int, optional):
+            number of channels in speaker embedding vectors. Defaults to 0.
+    """
+
+    num_chars: int = None
+    out_channels: int = 80
+    hidden_channels: int = 256
+    hidden_channels_dp: int = 256
+    encoder_type: str = "fftransformer"
+    encoder_params: dict = field(
+        default_factory=lambda: {"hidden_channels_ffn": 1024, "num_heads": 2, "num_layers": 6, "dropout_p": 0.1}
+    )
+    decoder_type: str = "fftransformer"
+    decoder_params: dict = field(
+        default_factory=lambda: {"hidden_channels_ffn": 1024, "num_heads": 2, "num_layers": 6, "dropout_p": 0.1}
+    )
+    length_scale: float = 1.0
+    num_speakers: int = 0
+    use_speaker_embedding: bool = False
+    use_d_vector_file: bool = False
+    d_vector_dim: int = 0
+
+
+class AlignTTS(BaseTTS):
+    """AlignTTS with modified duration predictor.
+    https://arxiv.org/pdf/2003.01950.pdf
+
+    Encoder -> DurationPredictor -> Decoder
+
+    Check :class:`AlignTTSArgs` for the class arguments.
+
+    Paper Abstract:
+        Targeting at both high efficiency and performance, we propose AlignTTS to predict the
+        mel-spectrum in parallel. AlignTTS is based on a Feed-Forward Transformer which generates mel-spectrum from a
+        sequence of characters, and the duration of each character is determined by a duration predictor.Instead of
+        adopting the attention mechanism in Transformer TTS to align text to mel-spectrum, the alignment loss is presented
+        to consider all possible alignments in training by use of dynamic programming. Experiments on the LJSpeech dataset s
+        how that our model achieves not only state-of-the-art performance which outperforms Transformer TTS by 0.03 in mean
+        option score (MOS), but also a high efficiency which is more than 50 times faster than real-time.
+
+    Note:
+        Original model uses a separate character embedding layer for duration predictor. However, it causes the
+        duration predictor to overfit and prevents learning higher level interactions among characters. Therefore,
+        we predict durations based on encoder outputs which has higher level information about input characters. This
+        enables training without phases as in the original paper.
+
+        Original model uses Transormers in encoder and decoder layers. However, here you can set the architecture
+        differently based on your requirements using ```encoder_type``` and ```decoder_type``` parameters.
+
+    Examples:
+        >>> from TTS.tts.configs.align_tts_config import AlignTTSConfig
+        >>> config = AlignTTSConfig()
+        >>> model = AlignTTS(config)
+
+    """
+
+    # pylint: disable=dangerous-default-value
+
+    def __init__(
+        self,
+        config: "AlignTTSConfig",
+        ap: "AudioProcessor" = None,
+        tokenizer: "TTSTokenizer" = None,
+        speaker_manager: SpeakerManager = None,
+    ):
+        super().__init__(config, ap, tokenizer, speaker_manager)
+        self.speaker_manager = speaker_manager
+        self.phase = -1
+        self.length_scale = (
+            float(config.model_args.length_scale)
+            if isinstance(config.model_args.length_scale, int)
+            else config.model_args.length_scale
+        )
+
+        self.emb = nn.Embedding(self.config.model_args.num_chars, self.config.model_args.hidden_channels)
+
+        self.embedded_speaker_dim = 0
+        self.init_multispeaker(config)
+
+        self.pos_encoder = PositionalEncoding(config.model_args.hidden_channels)
+        self.encoder = Encoder(
+            config.model_args.hidden_channels,
+            config.model_args.hidden_channels,
+            config.model_args.encoder_type,
+            config.model_args.encoder_params,
+            self.embedded_speaker_dim,
+        )
+        self.decoder = Decoder(
+            config.model_args.out_channels,
+            config.model_args.hidden_channels,
+            config.model_args.decoder_type,
+            config.model_args.decoder_params,
+        )
+        self.duration_predictor = DurationPredictor(config.model_args.hidden_channels_dp)
+
+        self.mod_layer = nn.Conv1d(config.model_args.hidden_channels, config.model_args.hidden_channels, 1)
+
+        self.mdn_block = MDNBlock(config.model_args.hidden_channels, 2 * config.model_args.out_channels)
+
+        if self.embedded_speaker_dim > 0 and self.embedded_speaker_dim != config.model_args.hidden_channels:
+            self.proj_g = nn.Conv1d(self.embedded_speaker_dim, config.model_args.hidden_channels, 1)
+
+    @staticmethod
+    def compute_log_probs(mu, log_sigma, y):
+        # pylint: disable=protected-access, c-extension-no-member
+        y = y.transpose(1, 2).unsqueeze(1)  # [B, 1, T1, D]
+        mu = mu.transpose(1, 2).unsqueeze(2)  # [B, T2, 1, D]
+        log_sigma = log_sigma.transpose(1, 2).unsqueeze(2)  # [B, T2, 1, D]
+        expanded_y, expanded_mu = torch.broadcast_tensors(y, mu)
+        exponential = -0.5 * torch.mean(
+            torch._C._nn.mse_loss(expanded_y, expanded_mu, 0) / torch.pow(log_sigma.exp(), 2), dim=-1
+        )  # B, L, T
+        logp = exponential - 0.5 * log_sigma.mean(dim=-1)
+        return logp
+
+    def compute_align_path(self, mu, log_sigma, y, x_mask, y_mask):
+        # find the max alignment path
+        attn_mask = torch.unsqueeze(x_mask, -1) * torch.unsqueeze(y_mask, 2)
+        log_p = self.compute_log_probs(mu, log_sigma, y)
+        # [B, T_en, T_dec]
+        attn = maximum_path(log_p, attn_mask.squeeze(1)).unsqueeze(1)
+        dr_mas = torch.sum(attn, -1)
+        return dr_mas.squeeze(1), log_p
+
+    @staticmethod
+    def generate_attn(dr, x_mask, y_mask=None):
+        # compute decode mask from the durations
+        if y_mask is None:
+            y_lengths = dr.sum(1).long()
+            y_lengths[y_lengths < 1] = 1
+            y_mask = torch.unsqueeze(sequence_mask(y_lengths, None), 1).to(dr.dtype)
+        attn_mask = torch.unsqueeze(x_mask, -1) * torch.unsqueeze(y_mask, 2)
+        attn = generate_path(dr, attn_mask.squeeze(1)).to(dr.dtype)
+        return attn
+
+    def expand_encoder_outputs(self, en, dr, x_mask, y_mask):
+        """Generate attention alignment map from durations and
+        expand encoder outputs
+
+        Examples::
+            - encoder output: [a,b,c,d]
+            - durations: [1, 3, 2, 1]
+
+            - expanded: [a, b, b, b, c, c, d]
+            - attention map: [[0, 0, 0, 0, 0, 0, 1],
+                             [0, 0, 0, 0, 1, 1, 0],
+                             [0, 1, 1, 1, 0, 0, 0],
+                             [1, 0, 0, 0, 0, 0, 0]]
+        """
+        attn = self.generate_attn(dr, x_mask, y_mask)
+        o_en_ex = torch.matmul(attn.squeeze(1).transpose(1, 2), en.transpose(1, 2)).transpose(1, 2)
+        return o_en_ex, attn
+
+    def format_durations(self, o_dr_log, x_mask):
+        o_dr = (torch.exp(o_dr_log) - 1) * x_mask * self.length_scale
+        o_dr[o_dr < 1] = 1.0
+        o_dr = torch.round(o_dr)
+        return o_dr
+
+    @staticmethod
+    def _concat_speaker_embedding(o_en, g):
+        g_exp = g.expand(-1, -1, o_en.size(-1))  # [B, C, T_en]
+        o_en = torch.cat([o_en, g_exp], 1)
+        return o_en
+
+    def _sum_speaker_embedding(self, x, g):
+        # project g to decoder dim.
+        if hasattr(self, "proj_g"):
+            g = self.proj_g(g)
+
+        return x + g
+
+    def _forward_encoder(self, x, x_lengths, g=None):
+        if hasattr(self, "emb_g"):
+            g = nn.functional.normalize(self.speaker_embedding(g))  # [B, C, 1]
+
+        if g is not None:
+            g = g.unsqueeze(-1)
+
+        # [B, T, C]
+        x_emb = self.emb(x)
+        # [B, C, T]
+        x_emb = torch.transpose(x_emb, 1, -1)
+
+        # compute sequence masks
+        x_mask = torch.unsqueeze(sequence_mask(x_lengths, x.shape[1]), 1).to(x.dtype)
+
+        # encoder pass
+        o_en = self.encoder(x_emb, x_mask)
+
+        # speaker conditioning for duration predictor
+        if g is not None:
+            o_en_dp = self._concat_speaker_embedding(o_en, g)
+        else:
+            o_en_dp = o_en
+        return o_en, o_en_dp, x_mask, g
+
+    def _forward_decoder(self, o_en, o_en_dp, dr, x_mask, y_lengths, g):
+        y_mask = torch.unsqueeze(sequence_mask(y_lengths, None), 1).to(o_en_dp.dtype)
+        # expand o_en with durations
+        o_en_ex, attn = self.expand_encoder_outputs(o_en, dr, x_mask, y_mask)
+        # positional encoding
+        if hasattr(self, "pos_encoder"):
+            o_en_ex = self.pos_encoder(o_en_ex, y_mask)
+        # speaker embedding
+        if g is not None:
+            o_en_ex = self._sum_speaker_embedding(o_en_ex, g)
+        # decoder pass
+        o_de = self.decoder(o_en_ex, y_mask, g=g)
+        return o_de, attn.transpose(1, 2)
+
+    def _forward_mdn(self, o_en, y, y_lengths, x_mask):
+        # MAS potentials and alignment
+        mu, log_sigma = self.mdn_block(o_en)
+        y_mask = torch.unsqueeze(sequence_mask(y_lengths, None), 1).to(o_en.dtype)
+        dr_mas, logp = self.compute_align_path(mu, log_sigma, y, x_mask, y_mask)
+        return dr_mas, mu, log_sigma, logp
+
+    def forward(
+        self, x, x_lengths, y, y_lengths, aux_input={"d_vectors": None}, phase=None
+    ):  # pylint: disable=unused-argument
+        """
+        Shapes:
+            - x: :math:`[B, T_max]`
+            - x_lengths: :math:`[B]`
+            - y_lengths: :math:`[B]`
+            - dr: :math:`[B, T_max]`
+            - g: :math:`[B, C]`
+        """
+        y = y.transpose(1, 2)
+        g = aux_input["d_vectors"] if "d_vectors" in aux_input else None
+        o_de, o_dr_log, dr_mas_log, attn, mu, log_sigma, logp = None, None, None, None, None, None, None
+        if phase == 0:
+            # train encoder and MDN
+            o_en, o_en_dp, x_mask, g = self._forward_encoder(x, x_lengths, g)
+            dr_mas, mu, log_sigma, logp = self._forward_mdn(o_en, y, y_lengths, x_mask)
+            y_mask = torch.unsqueeze(sequence_mask(y_lengths, None), 1).to(o_en_dp.dtype)
+            attn = self.generate_attn(dr_mas, x_mask, y_mask)
+        elif phase == 1:
+            # train decoder
+            o_en, o_en_dp, x_mask, g = self._forward_encoder(x, x_lengths, g)
+            dr_mas, _, _, _ = self._forward_mdn(o_en, y, y_lengths, x_mask)
+            o_de, attn = self._forward_decoder(o_en.detach(), o_en_dp.detach(), dr_mas.detach(), x_mask, y_lengths, g=g)
+        elif phase == 2:
+            # train the whole except duration predictor
+            o_en, o_en_dp, x_mask, g = self._forward_encoder(x, x_lengths, g)
+            dr_mas, mu, log_sigma, logp = self._forward_mdn(o_en, y, y_lengths, x_mask)
+            o_de, attn = self._forward_decoder(o_en, o_en_dp, dr_mas, x_mask, y_lengths, g=g)
+        elif phase == 3:
+            # train duration predictor
+            o_en, o_en_dp, x_mask, g = self._forward_encoder(x, x_lengths, g)
+            o_dr_log = self.duration_predictor(x, x_mask)
+            dr_mas, mu, log_sigma, logp = self._forward_mdn(o_en, y, y_lengths, x_mask)
+            o_de, attn = self._forward_decoder(o_en, o_en_dp, dr_mas, x_mask, y_lengths, g=g)
+            o_dr_log = o_dr_log.squeeze(1)
+        else:
+            o_en, o_en_dp, x_mask, g = self._forward_encoder(x, x_lengths, g)
+            o_dr_log = self.duration_predictor(o_en_dp.detach(), x_mask)
+            dr_mas, mu, log_sigma, logp = self._forward_mdn(o_en, y, y_lengths, x_mask)
+            o_de, attn = self._forward_decoder(o_en, o_en_dp, dr_mas, x_mask, y_lengths, g=g)
+            o_dr_log = o_dr_log.squeeze(1)
+        dr_mas_log = torch.log(dr_mas + 1).squeeze(1)
+        outputs = {
+            "model_outputs": o_de.transpose(1, 2),
+            "alignments": attn,
+            "durations_log": o_dr_log,
+            "durations_mas_log": dr_mas_log,
+            "mu": mu,
+            "log_sigma": log_sigma,
+            "logp": logp,
+        }
+        return outputs
+
+    @torch.no_grad()
+    def inference(self, x, aux_input={"d_vectors": None}):  # pylint: disable=unused-argument
+        """
+        Shapes:
+            - x: :math:`[B, T_max]`
+            - x_lengths: :math:`[B]`
+            - g: :math:`[B, C]`
+        """
+        g = aux_input["d_vectors"] if "d_vectors" in aux_input else None
+        x_lengths = torch.tensor(x.shape[1:2]).to(x.device)
+        # pad input to prevent dropping the last word
+        # x = torch.nn.functional.pad(x, pad=(0, 5), mode='constant', value=0)
+        o_en, o_en_dp, x_mask, g = self._forward_encoder(x, x_lengths, g)
+        # o_dr_log = self.duration_predictor(x, x_mask)
+        o_dr_log = self.duration_predictor(o_en_dp, x_mask)
+        # duration predictor pass
+        o_dr = self.format_durations(o_dr_log, x_mask).squeeze(1)
+        y_lengths = o_dr.sum(1)
+        o_de, attn = self._forward_decoder(o_en, o_en_dp, o_dr, x_mask, y_lengths, g=g)
+        outputs = {"model_outputs": o_de.transpose(1, 2), "alignments": attn}
+        return outputs
+
+    def train_step(self, batch: dict, criterion: nn.Module):
+        text_input = batch["text_input"]
+        text_lengths = batch["text_lengths"]
+        mel_input = batch["mel_input"]
+        mel_lengths = batch["mel_lengths"]
+        d_vectors = batch["d_vectors"]
+        speaker_ids = batch["speaker_ids"]
+
+        aux_input = {"d_vectors": d_vectors, "speaker_ids": speaker_ids}
+        outputs = self.forward(text_input, text_lengths, mel_input, mel_lengths, aux_input, self.phase)
+        loss_dict = criterion(
+            outputs["logp"],
+            outputs["model_outputs"],
+            mel_input,
+            mel_lengths,
+            outputs["durations_log"],
+            outputs["durations_mas_log"],
+            text_lengths,
+            phase=self.phase,
+        )
+
+        return outputs, loss_dict
+
+    def _create_logs(self, batch, outputs, ap):  # pylint: disable=no-self-use
+        model_outputs = outputs["model_outputs"]
+        alignments = outputs["alignments"]
+        mel_input = batch["mel_input"]
+
+        pred_spec = model_outputs[0].data.cpu().numpy()
+        gt_spec = mel_input[0].data.cpu().numpy()
+        align_img = alignments[0].data.cpu().numpy()
+
+        figures = {
+            "prediction": plot_spectrogram(pred_spec, ap, output_fig=False),
+            "ground_truth": plot_spectrogram(gt_spec, ap, output_fig=False),
+            "alignment": plot_alignment(align_img, output_fig=False),
+        }
+
+        # Sample audio
+        train_audio = ap.inv_melspectrogram(pred_spec.T)
+        return figures, {"audio": train_audio}
+
+    def train_log(
+        self, batch: dict, outputs: dict, logger: "Logger", assets: dict, steps: int
+    ) -> None:  # pylint: disable=no-self-use
+        figures, audios = self._create_logs(batch, outputs, self.ap)
+        logger.train_figures(steps, figures)
+        logger.train_audios(steps, audios, self.ap.sample_rate)
+
+    def eval_step(self, batch: dict, criterion: nn.Module):
+        return self.train_step(batch, criterion)
+
+    def eval_log(self, batch: dict, outputs: dict, logger: "Logger", assets: dict, steps: int) -> None:
+        figures, audios = self._create_logs(batch, outputs, self.ap)
+        logger.eval_figures(steps, figures)
+        logger.eval_audios(steps, audios, self.ap.sample_rate)
+
+    def load_checkpoint(
+        self, config, checkpoint_path, eval=False, cache=False
+    ):  # pylint: disable=unused-argument, redefined-builtin
+        state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"), cache=cache)
+        self.load_state_dict(state["model"])
+        if eval:
+            self.eval()
+            assert not self.training
+
+    def get_criterion(self):
+        from TTS.tts.layers.losses import AlignTTSLoss  # pylint: disable=import-outside-toplevel
+
+        return AlignTTSLoss(self.config)
+
+    @staticmethod
+    def _set_phase(config, global_step):
+        """Decide AlignTTS training phase"""
+        if isinstance(config.phase_start_steps, list):
+            vals = [i < global_step for i in config.phase_start_steps]
+            if not True in vals:
+                phase = 0
+            else:
+                phase = (
+                    len(config.phase_start_steps)
+                    - [i < global_step for i in config.phase_start_steps][::-1].index(True)
+                    - 1
+                )
+        else:
+            phase = None
+        return phase
+
+    def on_epoch_start(self, trainer):
+        """Set AlignTTS training phase on epoch start."""
+        self.phase = self._set_phase(trainer.config, trainer.total_steps_done)
+
+    @staticmethod
+    def init_from_config(config: "AlignTTSConfig", samples: Union[List[List], List[Dict]] = None):
+        """Initiate model from config
+
+        Args:
+            config (AlignTTSConfig): Model config.
+            samples (Union[List[List], List[Dict]]): Training samples to parse speaker ids for training.
+                Defaults to None.
+        """
+        from TTS.utils.audio import AudioProcessor
+
+        ap = AudioProcessor.init_from_config(config)
+        tokenizer, new_config = TTSTokenizer.init_from_config(config)
+        speaker_manager = SpeakerManager.init_from_config(config, samples)
+        return AlignTTS(new_config, ap, tokenizer, speaker_manager)

TTS/tts/models/bark.py

@@ -0,0 +1,284 @@
+import os
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+from coqpit import Coqpit
+from encodec import EncodecModel
+from transformers import BertTokenizer
+
+from TTS.tts.layers.bark.inference_funcs import (
+    codec_decode,
+    generate_coarse,
+    generate_fine,
+    generate_text_semantic,
+    generate_voice,
+    load_voice,
+)
+from TTS.tts.layers.bark.load_model import load_model
+from TTS.tts.layers.bark.model import GPT
+from TTS.tts.layers.bark.model_fine import FineGPT
+from TTS.tts.models.base_tts import BaseTTS
+
+
+@dataclass
+class BarkAudioConfig(Coqpit):
+    sample_rate: int = 24000
+    output_sample_rate: int = 24000
+
+
+class Bark(BaseTTS):
+    def __init__(
+        self,
+        config: Coqpit,
+        tokenizer: BertTokenizer = BertTokenizer.from_pretrained("bert-base-multilingual-cased"),
+    ) -> None:
+        super().__init__(config=config, ap=None, tokenizer=None, speaker_manager=None, language_manager=None)
+        self.config.num_chars = len(tokenizer)
+        self.tokenizer = tokenizer
+        self.semantic_model = GPT(config.semantic_config)
+        self.coarse_model = GPT(config.coarse_config)
+        self.fine_model = FineGPT(config.fine_config)
+        self.encodec = EncodecModel.encodec_model_24khz()
+        self.encodec.set_target_bandwidth(6.0)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def load_bark_models(self):
+        self.semantic_model, self.config = load_model(
+            ckpt_path=self.config.LOCAL_MODEL_PATHS["text"], device=self.device, config=self.config, model_type="text"
+        )
+        self.coarse_model, self.config = load_model(
+            ckpt_path=self.config.LOCAL_MODEL_PATHS["coarse"],
+            device=self.device,
+            config=self.config,
+            model_type="coarse",
+        )
+        self.fine_model, self.config = load_model(
+            ckpt_path=self.config.LOCAL_MODEL_PATHS["fine"], device=self.device, config=self.config, model_type="fine"
+        )
+
+    def train_step(
+        self,
+    ):
+        pass
+
+    def text_to_semantic(
+        self,
+        text: str,
+        history_prompt: Optional[str] = None,
+        temp: float = 0.7,
+        base=None,
+        allow_early_stop=True,
+        **kwargs,
+    ):
+        """Generate semantic array from text.
+
+        Args:
+            text: text to be turned into audio
+            history_prompt: history choice for audio cloning
+            temp: generation temperature (1.0 more diverse, 0.0 more conservative)
+
+        Returns:
+            numpy semantic array to be fed into `semantic_to_waveform`
+        """
+        x_semantic = generate_text_semantic(
+            text,
+            self,
+            history_prompt=history_prompt,
+            temp=temp,
+            base=base,
+            allow_early_stop=allow_early_stop,
+            **kwargs,
+        )
+        return x_semantic
+
+    def semantic_to_waveform(
+        self,
+        semantic_tokens: np.ndarray,
+        history_prompt: Optional[str] = None,
+        temp: float = 0.7,
+        base=None,
+    ):
+        """Generate audio array from semantic input.
+
+        Args:
+            semantic_tokens: semantic token output from `text_to_semantic`
+            history_prompt: history choice for audio cloning
+            temp: generation temperature (1.0 more diverse, 0.0 more conservative)
+
+        Returns:
+            numpy audio array at sample frequency 24khz
+        """
+        x_coarse_gen = generate_coarse(
+            semantic_tokens,
+            self,
+            history_prompt=history_prompt,
+            temp=temp,
+            base=base,
+        )
+        x_fine_gen = generate_fine(
+            x_coarse_gen,
+            self,
+            history_prompt=history_prompt,
+            temp=0.5,
+            base=base,
+        )
+        audio_arr = codec_decode(x_fine_gen, self)
+        return audio_arr, x_coarse_gen, x_fine_gen
+
+    def generate_audio(
+        self,
+        text: str,
+        history_prompt: Optional[str] = None,
+        text_temp: float = 0.7,
+        waveform_temp: float = 0.7,
+        base=None,
+        allow_early_stop=True,
+        **kwargs,
+    ):
+        """Generate audio array from input text.
+
+        Args:
+            text: text to be turned into audio
+            history_prompt: history choice for audio cloning
+            text_temp: generation temperature (1.0 more diverse, 0.0 more conservative)
+            waveform_temp: generation temperature (1.0 more diverse, 0.0 more conservative)
+
+        Returns:
+            numpy audio array at sample frequency 24khz
+        """
+        x_semantic = self.text_to_semantic(
+            text,
+            history_prompt=history_prompt,
+            temp=text_temp,
+            base=base,
+            allow_early_stop=allow_early_stop,
+            **kwargs,
+        )
+        audio_arr, c, f = self.semantic_to_waveform(
+            x_semantic, history_prompt=history_prompt, temp=waveform_temp, base=base
+        )
+        return audio_arr, [x_semantic, c, f]
+
+    def generate_voice(self, audio, speaker_id, voice_dir):
+        """Generate a voice from the given audio and text.
+
+        Args:
+            audio (str): Path to the audio file.
+            speaker_id (str): Speaker name.
+            voice_dir (str): Path to the directory to save the generate voice.
+        """
+        if voice_dir is not None:
+            voice_dirs = [voice_dir]
+            try:
+                _ = load_voice(speaker_id, voice_dirs)
+            except (KeyError, FileNotFoundError):
+                output_path = os.path.join(voice_dir, speaker_id + ".npz")
+                os.makedirs(voice_dir, exist_ok=True)
+                generate_voice(audio, self, output_path)
+
+    def _set_voice_dirs(self, voice_dirs):
+        def_voice_dir = None
+        if isinstance(self.config.DEF_SPEAKER_DIR, str):
+            os.makedirs(self.config.DEF_SPEAKER_DIR, exist_ok=True)
+            if os.path.isdir(self.config.DEF_SPEAKER_DIR):
+                def_voice_dir = self.config.DEF_SPEAKER_DIR
+        _voice_dirs = [def_voice_dir] if def_voice_dir is not None else []
+        if voice_dirs is not None:
+            if isinstance(voice_dirs, str):
+                voice_dirs = [voice_dirs]
+            _voice_dirs = voice_dirs + _voice_dirs
+        return _voice_dirs
+
+    # TODO: remove config from synthesize
+    def synthesize(
+        self, text, config, speaker_id="random", voice_dirs=None, **kwargs
+    ):  # pylint: disable=unused-argument
+        """Synthesize speech with the given input text.
+
+        Args:
+            text (str): Input text.
+            config (BarkConfig): Config with inference parameters.
+            speaker_id (str): One of the available speaker names. If `random`, it generates a random speaker.
+            speaker_wav (str): Path to the speaker audio file for cloning a new voice. It is cloned and saved in
+                `voice_dirs` with the name `speaker_id`. Defaults to None.
+            voice_dirs (List[str]): List of paths that host reference audio files for speakers. Defaults to None.
+            **kwargs: Model specific inference settings used by `generate_audio()` and `TTS.tts.layers.bark.inference_funcs.generate_text_semantic().
+
+        Returns:
+            A dictionary of the output values with `wav` as output waveform, `deterministic_seed` as seed used at inference,
+            `text_input` as text token IDs after tokenizer, `voice_samples` as samples used for cloning, `conditioning_latents`
+            as latents used at inference.
+
+        """
+        speaker_id = "random" if speaker_id is None else speaker_id
+        voice_dirs = self._set_voice_dirs(voice_dirs)
+        history_prompt = load_voice(self, speaker_id, voice_dirs)
+        outputs = self.generate_audio(text, history_prompt=history_prompt, **kwargs)
+        return_dict = {
+            "wav": outputs[0],
+            "text_inputs": text,
+        }
+
+        return return_dict
+
+    def eval_step(self):
+        ...
+
+    def forward(self):
+        ...
+
+    def inference(self):
+        ...
+
+    @staticmethod
+    def init_from_config(config: "BarkConfig", **kwargs):  # pylint: disable=unused-argument
+        return Bark(config)
+
+    # pylint: disable=unused-argument, redefined-builtin
+    def load_checkpoint(
+        self,
+        config,
+        checkpoint_dir,
+        text_model_path=None,
+        coarse_model_path=None,
+        fine_model_path=None,
+        hubert_model_path=None,
+        hubert_tokenizer_path=None,
+        eval=False,
+        strict=True,
+        **kwargs,
+    ):
+        """Load a model checkpoints from a directory. This model is with multiple checkpoint files and it
+        expects to have all the files to be under the given `checkpoint_dir` with the rigth names.
+        If eval is True, set the model to eval mode.
+
+        Args:
+            config (TortoiseConfig): The model config.
+            checkpoint_dir (str): The directory where the checkpoints are stored.
+            ar_checkpoint_path (str, optional): The path to the autoregressive checkpoint. Defaults to None.
+            diff_checkpoint_path (str, optional): The path to the diffusion checkpoint. Defaults to None.
+            clvp_checkpoint_path (str, optional): The path to the CLVP checkpoint. Defaults to None.
+            vocoder_checkpoint_path (str, optional): The path to the vocoder checkpoint. Defaults to None.
+            eval (bool, optional): Whether to set the model to eval mode. Defaults to False.
+            strict (bool, optional): Whether to load the model strictly. Defaults to True.
+        """
+        text_model_path = text_model_path or os.path.join(checkpoint_dir, "text_2.pt")
+        coarse_model_path = coarse_model_path or os.path.join(checkpoint_dir, "coarse_2.pt")
+        fine_model_path = fine_model_path or os.path.join(checkpoint_dir, "fine_2.pt")
+        hubert_model_path = hubert_model_path or os.path.join(checkpoint_dir, "hubert.pt")
+        hubert_tokenizer_path = hubert_tokenizer_path or os.path.join(checkpoint_dir, "tokenizer.pth")
+
+        self.config.LOCAL_MODEL_PATHS["text"] = text_model_path
+        self.config.LOCAL_MODEL_PATHS["coarse"] = coarse_model_path
+        self.config.LOCAL_MODEL_PATHS["fine"] = fine_model_path
+        self.config.LOCAL_MODEL_PATHS["hubert"] = hubert_model_path
+        self.config.LOCAL_MODEL_PATHS["hubert_tokenizer"] = hubert_tokenizer_path
+
+        self.load_bark_models()
+
+        if eval:
+            self.eval()

TTS/tts/models/base_tacotron.py

@@ -0,0 +1,305 @@
+import copy
+from abc import abstractmethod
+from typing import Dict, Tuple
+
+import torch
+from coqpit import Coqpit
+from torch import nn
+
+from TTS.tts.layers.losses import TacotronLoss
+from TTS.tts.models.base_tts import BaseTTS
+from TTS.tts.utils.helpers import sequence_mask
+from TTS.tts.utils.speakers import SpeakerManager
+from TTS.tts.utils.synthesis import synthesis
+from TTS.tts.utils.text.tokenizer import TTSTokenizer
+from TTS.tts.utils.visual import plot_alignment, plot_spectrogram
+from TTS.utils.generic_utils import format_aux_input
+from TTS.utils.io import load_fsspec
+from TTS.utils.training import gradual_training_scheduler
+
+
+class BaseTacotron(BaseTTS):
+    """Base class shared by Tacotron and Tacotron2"""
+
+    def __init__(
+        self,
+        config: "TacotronConfig",
+        ap: "AudioProcessor",
+        tokenizer: "TTSTokenizer",
+        speaker_manager: SpeakerManager = None,
+    ):
+        super().__init__(config, ap, tokenizer, speaker_manager)
+
+        # pass all config fields as class attributes
+        for key in config:
+            setattr(self, key, config[key])
+
+        # layers
+        self.embedding = None
+        self.encoder = None
+        self.decoder = None
+        self.postnet = None
+
+        # init tensors
+        self.embedded_speakers = None
+        self.embedded_speakers_projected = None
+
+        # global style token
+        if self.gst and self.use_gst:
+            self.decoder_in_features += self.gst.gst_embedding_dim  # add gst embedding dim
+            self.gst_layer = None
+
+        # Capacitron
+        if self.capacitron_vae and self.use_capacitron_vae:
+            self.decoder_in_features += self.capacitron_vae.capacitron_VAE_embedding_dim  # add capacitron embedding dim
+            self.capacitron_vae_layer = None
+
+        # additional layers
+        self.decoder_backward = None
+        self.coarse_decoder = None
+
+    @staticmethod
+    def _format_aux_input(aux_input: Dict) -> Dict:
+        """Set missing fields to their default values"""
+        if aux_input:
+            return format_aux_input({"d_vectors": None, "speaker_ids": None}, aux_input)
+        return None
+
+    #############################
+    # INIT FUNCTIONS
+    #############################
+
+    def _init_backward_decoder(self):
+        """Init the backward decoder for Forward-Backward decoding."""
+        self.decoder_backward = copy.deepcopy(self.decoder)
+
+    def _init_coarse_decoder(self):
+        """Init the coarse decoder for Double-Decoder Consistency."""
+        self.coarse_decoder = copy.deepcopy(self.decoder)
+        self.coarse_decoder.r_init = self.ddc_r
+        self.coarse_decoder.set_r(self.ddc_r)
+
+    #############################
+    # CORE FUNCTIONS
+    #############################
+
+    @abstractmethod
+    def forward(self):
+        pass
+
+    @abstractmethod
+    def inference(self):
+        pass
+
+    def load_checkpoint(
+        self, config, checkpoint_path, eval=False, cache=False
+    ):  # pylint: disable=unused-argument, redefined-builtin
+        """Load model checkpoint and set up internals.
+
+        Args:
+            config (Coqpi): model configuration.
+            checkpoint_path (str): path to checkpoint file.
+            eval (bool, optional): whether to load model for evaluation.
+            cache (bool, optional): If True, cache the file locally for subsequent calls. It is cached under `get_user_data_dir()/tts_cache`. Defaults to False.
+        """
+        state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"), cache=cache)
+        self.load_state_dict(state["model"])
+        # TODO: set r in run-time by taking it from the new config
+        if "r" in state:
+            # set r from the state (for compatibility with older checkpoints)
+            self.decoder.set_r(state["r"])
+        elif "config" in state:
+            # set r from config used at training time (for inference)
+            self.decoder.set_r(state["config"]["r"])
+        else:
+            # set r from the new config (for new-models)
+            self.decoder.set_r(config.r)
+        if eval:
+            self.eval()
+            print(f" > Model's reduction rate `r` is set to: {self.decoder.r}")
+            assert not self.training
+
+    def get_criterion(self) -> nn.Module:
+        """Get the model criterion used in training."""
+        return TacotronLoss(self.config)
+
+    @staticmethod
+    def init_from_config(config: Coqpit):
+        """Initialize model from config."""
+        from TTS.utils.audio import AudioProcessor
+
+        ap = AudioProcessor.init_from_config(config)
+        tokenizer = TTSTokenizer.init_from_config(config)
+        speaker_manager = SpeakerManager.init_from_config(config)
+        return BaseTacotron(config, ap, tokenizer, speaker_manager)
+
+    ##########################
+    # TEST AND LOG FUNCTIONS #
+    ##########################
+
+    def test_run(self, assets: Dict) -> Tuple[Dict, Dict]:
+        """Generic test run for `tts` models used by `Trainer`.
+
+        You can override this for a different behaviour.
+
+        Args:
+            assets (dict): A dict of training assets. For `tts` models, it must include `{'audio_processor': ap}`.
+
+        Returns:
+            Tuple[Dict, Dict]: Test figures and audios to be projected to Tensorboard.
+        """
+        print(" | > Synthesizing test sentences.")
+        test_audios = {}
+        test_figures = {}
+        test_sentences = self.config.test_sentences
+        aux_inputs = self._get_test_aux_input()
+        for idx, sen in enumerate(test_sentences):
+            outputs_dict = synthesis(
+                self,
+                sen,
+                self.config,
+                "cuda" in str(next(self.parameters()).device),
+                speaker_id=aux_inputs["speaker_id"],
+                d_vector=aux_inputs["d_vector"],
+                style_wav=aux_inputs["style_wav"],
+                use_griffin_lim=True,
+                do_trim_silence=False,
+            )
+            test_audios["{}-audio".format(idx)] = outputs_dict["wav"]
+            test_figures["{}-prediction".format(idx)] = plot_spectrogram(
+                outputs_dict["outputs"]["model_outputs"], self.ap, output_fig=False
+            )
+            test_figures["{}-alignment".format(idx)] = plot_alignment(
+                outputs_dict["outputs"]["alignments"], output_fig=False
+            )
+        return {"figures": test_figures, "audios": test_audios}
+
+    def test_log(
+        self, outputs: dict, logger: "Logger", assets: dict, steps: int  # pylint: disable=unused-argument
+    ) -> None:
+        logger.test_audios(steps, outputs["audios"], self.ap.sample_rate)
+        logger.test_figures(steps, outputs["figures"])
+
+    #############################
+    # COMMON COMPUTE FUNCTIONS
+    #############################
+
+    def compute_masks(self, text_lengths, mel_lengths):
+        """Compute masks  against sequence paddings."""
+        # B x T_in_max (boolean)
+        input_mask = sequence_mask(text_lengths)
+        output_mask = None
+        if mel_lengths is not None:
+            max_len = mel_lengths.max()
+            r = self.decoder.r
+            max_len = max_len + (r - (max_len % r)) if max_len % r > 0 else max_len
+            output_mask = sequence_mask(mel_lengths, max_len=max_len)
+        return input_mask, output_mask
+
+    def _backward_pass(self, mel_specs, encoder_outputs, mask):
+        """Run backwards decoder"""
+        decoder_outputs_b, alignments_b, _ = self.decoder_backward(
+            encoder_outputs, torch.flip(mel_specs, dims=(1,)), mask
+        )
+        decoder_outputs_b = decoder_outputs_b.transpose(1, 2).contiguous()
+        return decoder_outputs_b, alignments_b
+
+    def _coarse_decoder_pass(self, mel_specs, encoder_outputs, alignments, input_mask):
+        """Double Decoder Consistency"""
+        T = mel_specs.shape[1]
+        if T % self.coarse_decoder.r > 0:
+            padding_size = self.coarse_decoder.r - (T % self.coarse_decoder.r)
+            mel_specs = torch.nn.functional.pad(mel_specs, (0, 0, 0, padding_size, 0, 0))
+        decoder_outputs_backward, alignments_backward, _ = self.coarse_decoder(
+            encoder_outputs.detach(), mel_specs, input_mask
+        )
+        # scale_factor = self.decoder.r_init / self.decoder.r
+        alignments_backward = torch.nn.functional.interpolate(
+            alignments_backward.transpose(1, 2),
+            size=alignments.shape[1],
+            mode="nearest",
+        ).transpose(1, 2)
+        decoder_outputs_backward = decoder_outputs_backward.transpose(1, 2)
+        decoder_outputs_backward = decoder_outputs_backward[:, :T, :]
+        return decoder_outputs_backward, alignments_backward
+
+    #############################
+    # EMBEDDING FUNCTIONS
+    #############################
+
+    def compute_gst(self, inputs, style_input, speaker_embedding=None):
+        """Compute global style token"""
+        if isinstance(style_input, dict):
+            # multiply each style token with a weight
+            query = torch.zeros(1, 1, self.gst.gst_embedding_dim // 2).type_as(inputs)
+            if speaker_embedding is not None:
+                query = torch.cat([query, speaker_embedding.reshape(1, 1, -1)], dim=-1)
+
+            _GST = torch.tanh(self.gst_layer.style_token_layer.style_tokens)
+            gst_outputs = torch.zeros(1, 1, self.gst.gst_embedding_dim).type_as(inputs)
+            for k_token, v_amplifier in style_input.items():
+                key = _GST[int(k_token)].unsqueeze(0).expand(1, -1, -1)
+                gst_outputs_att = self.gst_layer.style_token_layer.attention(query, key)
+                gst_outputs = gst_outputs + gst_outputs_att * v_amplifier
+        elif style_input is None:
+            # ignore style token and return zero tensor
+            gst_outputs = torch.zeros(1, 1, self.gst.gst_embedding_dim).type_as(inputs)
+        else:
+            # compute style tokens
+            gst_outputs = self.gst_layer(style_input, speaker_embedding)  # pylint: disable=not-callable
+        inputs = self._concat_speaker_embedding(inputs, gst_outputs)
+        return inputs
+
+    def compute_capacitron_VAE_embedding(self, inputs, reference_mel_info, text_info=None, speaker_embedding=None):
+        """Capacitron Variational Autoencoder"""
+        (
+            VAE_outputs,
+            posterior_distribution,
+            prior_distribution,
+            capacitron_beta,
+        ) = self.capacitron_vae_layer(
+            reference_mel_info,
+            text_info,
+            speaker_embedding,  # pylint: disable=not-callable
+        )
+
+        VAE_outputs = VAE_outputs.to(inputs.device)
+        encoder_output = self._concat_speaker_embedding(
+            inputs, VAE_outputs
+        )  # concatenate to the output of the basic tacotron encoder
+        return (
+            encoder_output,
+            posterior_distribution,
+            prior_distribution,
+            capacitron_beta,
+        )
+
+    @staticmethod
+    def _add_speaker_embedding(outputs, embedded_speakers):
+        embedded_speakers_ = embedded_speakers.expand(outputs.size(0), outputs.size(1), -1)
+        outputs = outputs + embedded_speakers_
+        return outputs
+
+    @staticmethod
+    def _concat_speaker_embedding(outputs, embedded_speakers):
+        embedded_speakers_ = embedded_speakers.expand(outputs.size(0), outputs.size(1), -1)
+        outputs = torch.cat([outputs, embedded_speakers_], dim=-1)
+        return outputs
+
+    #############################
+    # CALLBACKS
+    #############################
+
+    def on_epoch_start(self, trainer):
+        """Callback for setting values wrt gradual training schedule.
+
+        Args:
+            trainer (TrainerTTS): TTS trainer object that is used to train this model.
+        """
+        if self.gradual_training:
+            r, trainer.config.batch_size = gradual_training_scheduler(trainer.total_steps_done, trainer.config)
+            trainer.config.r = r
+            self.decoder.set_r(r)
+            if trainer.config.bidirectional_decoder:
+                trainer.model.decoder_backward.set_r(r)
+            print(f"\n > Number of output frames: {self.decoder.r}")

TTS/tts/models/base_tts.py

@@ -0,0 +1,459 @@
+import os
+import random
+from typing import Dict, List, Tuple, Union
+
+import torch
+import torch.distributed as dist
+from coqpit import Coqpit
+from torch import nn
+from torch.utils.data import DataLoader
+from torch.utils.data.sampler import WeightedRandomSampler
+from trainer.torch import DistributedSampler, DistributedSamplerWrapper
+
+from TTS.model import BaseTrainerModel
+from TTS.tts.datasets.dataset import TTSDataset
+from TTS.tts.utils.data import get_length_balancer_weights
+from TTS.tts.utils.languages import LanguageManager, get_language_balancer_weights
+from TTS.tts.utils.speakers import SpeakerManager, get_speaker_balancer_weights, get_speaker_manager
+from TTS.tts.utils.synthesis import synthesis
+from TTS.tts.utils.visual import plot_alignment, plot_spectrogram
+
+# pylint: skip-file
+
+
+class BaseTTS(BaseTrainerModel):
+    """Base `tts` class. Every new `tts` model must inherit this.
+
+    It defines common `tts` specific functions on top of `Model` implementation.
+    """
+
+    MODEL_TYPE = "tts"
+
+    def __init__(
+        self,
+        config: Coqpit,
+        ap: "AudioProcessor",
+        tokenizer: "TTSTokenizer",
+        speaker_manager: SpeakerManager = None,
+        language_manager: LanguageManager = None,
+    ):
+        super().__init__()
+        self.config = config
+        self.ap = ap
+        self.tokenizer = tokenizer
+        self.speaker_manager = speaker_manager
+        self.language_manager = language_manager
+        self._set_model_args(config)
+
+    def _set_model_args(self, config: Coqpit):
+        """Setup model args based on the config type (`ModelConfig` or `ModelArgs`).
+
+        `ModelArgs` has all the fields reuqired to initialize the model architecture.
+
+        `ModelConfig` has all the fields required for training, inference and containes `ModelArgs`.
+
+        If the config is for training with a name like "*Config", then the model args are embeded in the
+        config.model_args
+
+        If the config is for the model with a name like "*Args", then we assign the directly.
+        """
+        # don't use isintance not to import recursively
+        if "Config" in config.__class__.__name__:
+            config_num_chars = (
+                self.config.model_args.num_chars if hasattr(self.config, "model_args") else self.config.num_chars
+            )
+            num_chars = config_num_chars if self.tokenizer is None else self.tokenizer.characters.num_chars
+            if "characters" in config:
+                self.config.num_chars = num_chars
+                if hasattr(self.config, "model_args"):
+                    config.model_args.num_chars = num_chars
+                    self.args = self.config.model_args
+            else:
+                self.config = config
+                self.args = config.model_args
+        elif "Args" in config.__class__.__name__:
+            self.args = config
+        else:
+            raise ValueError("config must be either a *Config or *Args")
+
+    def init_multispeaker(self, config: Coqpit, data: List = None):
+        """Initialize a speaker embedding layer if needen and define expected embedding channel size for defining
+        `in_channels` size of the connected layers.
+
+        This implementation yields 3 possible outcomes:
+
+        1. If `config.use_speaker_embedding` and `config.use_d_vector_file are False, do nothing.
+        2. If `config.use_d_vector_file` is True, set expected embedding channel size to `config.d_vector_dim` or 512.
+        3. If `config.use_speaker_embedding`, initialize a speaker embedding layer with channel size of
+        `config.d_vector_dim` or 512.
+
+        You can override this function for new models.
+
+        Args:
+            config (Coqpit): Model configuration.
+        """
+        # set number of speakers
+        if self.speaker_manager is not None:
+            self.num_speakers = self.speaker_manager.num_speakers
+        elif hasattr(config, "num_speakers"):
+            self.num_speakers = config.num_speakers
+
+        # set ultimate speaker embedding size
+        if config.use_speaker_embedding or config.use_d_vector_file:
+            self.embedded_speaker_dim = (
+                config.d_vector_dim if "d_vector_dim" in config and config.d_vector_dim is not None else 512
+            )
+        # init speaker embedding layer
+        if config.use_speaker_embedding and not config.use_d_vector_file:
+            print(" > Init speaker_embedding layer.")
+            self.speaker_embedding = nn.Embedding(self.num_speakers, self.embedded_speaker_dim)
+            self.speaker_embedding.weight.data.normal_(0, 0.3)
+
+    def get_aux_input(self, **kwargs) -> Dict:
+        """Prepare and return `aux_input` used by `forward()`"""
+        return {"speaker_id": None, "style_wav": None, "d_vector": None, "language_id": None}
+
+    def get_aux_input_from_test_sentences(self, sentence_info):
+        if hasattr(self.config, "model_args"):
+            config = self.config.model_args
+        else:
+            config = self.config
+
+        # extract speaker and language info
+        text, speaker_name, style_wav, language_name = None, None, None, None
+
+        if isinstance(sentence_info, list):
+            if len(sentence_info) == 1:
+                text = sentence_info[0]
+            elif len(sentence_info) == 2:
+                text, speaker_name = sentence_info
+            elif len(sentence_info) == 3:
+                text, speaker_name, style_wav = sentence_info
+            elif len(sentence_info) == 4:
+                text, speaker_name, style_wav, language_name = sentence_info
+        else:
+            text = sentence_info
+
+        # get speaker  id/d_vector
+        speaker_id, d_vector, language_id = None, None, None
+        if self.speaker_manager is not None:
+            if config.use_d_vector_file:
+                if speaker_name is None:
+                    d_vector = self.speaker_manager.get_random_embedding()
+                else:
+                    d_vector = self.speaker_manager.get_d_vector_by_name(speaker_name)
+            elif config.use_speaker_embedding:
+                if speaker_name is None:
+                    speaker_id = self.speaker_manager.get_random_id()
+                else:
+                    speaker_id = self.speaker_manager.name_to_id[speaker_name]
+
+        # get language id
+        if self.language_manager is not None and config.use_language_embedding and language_name is not None:
+            language_id = self.language_manager.name_to_id[language_name]
+
+        return {
+            "text": text,
+            "speaker_id": speaker_id,
+            "style_wav": style_wav,
+            "d_vector": d_vector,
+            "language_id": language_id,
+        }
+
+    def format_batch(self, batch: Dict) -> Dict:
+        """Generic batch formatting for `TTSDataset`.
+
+        You must override this if you use a custom dataset.
+
+        Args:
+            batch (Dict): [description]
+
+        Returns:
+            Dict: [description]
+        """
+        # setup input batch
+        text_input = batch["token_id"]
+        text_lengths = batch["token_id_lengths"]
+        speaker_names = batch["speaker_names"]
+        linear_input = batch["linear"]
+        mel_input = batch["mel"]
+        mel_lengths = batch["mel_lengths"]
+        stop_targets = batch["stop_targets"]
+        item_idx = batch["item_idxs"]
+        d_vectors = batch["d_vectors"]
+        speaker_ids = batch["speaker_ids"]
+        attn_mask = batch["attns"]
+        waveform = batch["waveform"]
+        pitch = batch["pitch"]
+        energy = batch["energy"]
+        language_ids = batch["language_ids"]
+        max_text_length = torch.max(text_lengths.float())
+        max_spec_length = torch.max(mel_lengths.float())
+
+        # compute durations from attention masks
+        durations = None
+        if attn_mask is not None:
+            durations = torch.zeros(attn_mask.shape[0], attn_mask.shape[2])
+            for idx, am in enumerate(attn_mask):
+                # compute raw durations
+                c_idxs = am[:, : text_lengths[idx], : mel_lengths[idx]].max(1)[1]
+                # c_idxs, counts = torch.unique_consecutive(c_idxs, return_counts=True)
+                c_idxs, counts = torch.unique(c_idxs, return_counts=True)
+                dur = torch.ones([text_lengths[idx]]).to(counts.dtype)
+                dur[c_idxs] = counts
+                # smooth the durations and set any 0 duration to 1
+                # by cutting off from the largest duration indeces.
+                extra_frames = dur.sum() - mel_lengths[idx]
+                largest_idxs = torch.argsort(-dur)[:extra_frames]
+                dur[largest_idxs] -= 1
+                assert (
+                    dur.sum() == mel_lengths[idx]
+                ), f" [!] total duration {dur.sum()} vs spectrogram length {mel_lengths[idx]}"
+                durations[idx, : text_lengths[idx]] = dur
+
+        # set stop targets wrt reduction factor
+        stop_targets = stop_targets.view(text_input.shape[0], stop_targets.size(1) // self.config.r, -1)
+        stop_targets = (stop_targets.sum(2) > 0.0).unsqueeze(2).float().squeeze(2)
+        stop_target_lengths = torch.divide(mel_lengths, self.config.r).ceil_()
+
+        return {
+            "text_input": text_input,
+            "text_lengths": text_lengths,
+            "speaker_names": speaker_names,
+            "mel_input": mel_input,
+            "mel_lengths": mel_lengths,
+            "linear_input": linear_input,
+            "stop_targets": stop_targets,
+            "stop_target_lengths": stop_target_lengths,
+            "attn_mask": attn_mask,
+            "durations": durations,
+            "speaker_ids": speaker_ids,
+            "d_vectors": d_vectors,
+            "max_text_length": float(max_text_length),
+            "max_spec_length": float(max_spec_length),
+            "item_idx": item_idx,
+            "waveform": waveform,
+            "pitch": pitch,
+            "energy": energy,
+            "language_ids": language_ids,
+            "audio_unique_names": batch["audio_unique_names"],
+        }
+
+    def get_sampler(self, config: Coqpit, dataset: TTSDataset, num_gpus=1):
+        weights = None
+        data_items = dataset.samples
+
+        if getattr(config, "use_language_weighted_sampler", False):
+            alpha = getattr(config, "language_weighted_sampler_alpha", 1.0)
+            print(" > Using Language weighted sampler with alpha:", alpha)
+            weights = get_language_balancer_weights(data_items) * alpha
+
+        if getattr(config, "use_speaker_weighted_sampler", False):
+            alpha = getattr(config, "speaker_weighted_sampler_alpha", 1.0)
+            print(" > Using Speaker weighted sampler with alpha:", alpha)
+            if weights is not None:
+                weights += get_speaker_balancer_weights(data_items) * alpha
+            else:
+                weights = get_speaker_balancer_weights(data_items) * alpha
+
+        if getattr(config, "use_length_weighted_sampler", False):
+            alpha = getattr(config, "length_weighted_sampler_alpha", 1.0)
+            print(" > Using Length weighted sampler with alpha:", alpha)
+            if weights is not None:
+                weights += get_length_balancer_weights(data_items) * alpha
+            else:
+                weights = get_length_balancer_weights(data_items) * alpha
+
+        if weights is not None:
+            sampler = WeightedRandomSampler(weights, len(weights))
+        else:
+            sampler = None
+
+        # sampler for DDP
+        if sampler is None:
+            sampler = DistributedSampler(dataset) if num_gpus > 1 else None
+        else:  # If a sampler is already defined use this sampler and DDP sampler together
+            sampler = DistributedSamplerWrapper(sampler) if num_gpus > 1 else sampler
+
+        return sampler
+
+    def get_data_loader(
+        self,
+        config: Coqpit,
+        assets: Dict,
+        is_eval: bool,
+        samples: Union[List[Dict], List[List]],
+        verbose: bool,
+        num_gpus: int,
+        rank: int = None,
+    ) -> "DataLoader":
+        if is_eval and not config.run_eval:
+            loader = None
+        else:
+            # setup multi-speaker attributes
+            if self.speaker_manager is not None:
+                if hasattr(config, "model_args"):
+                    speaker_id_mapping = (
+                        self.speaker_manager.name_to_id if config.model_args.use_speaker_embedding else None
+                    )
+                    d_vector_mapping = self.speaker_manager.embeddings if config.model_args.use_d_vector_file else None
+                    config.use_d_vector_file = config.model_args.use_d_vector_file
+                else:
+                    speaker_id_mapping = self.speaker_manager.name_to_id if config.use_speaker_embedding else None
+                    d_vector_mapping = self.speaker_manager.embeddings if config.use_d_vector_file else None
+            else:
+                speaker_id_mapping = None
+                d_vector_mapping = None
+
+            # setup multi-lingual attributes
+            if self.language_manager is not None:
+                language_id_mapping = self.language_manager.name_to_id if self.args.use_language_embedding else None
+            else:
+                language_id_mapping = None
+
+            # init dataloader
+            dataset = TTSDataset(
+                outputs_per_step=config.r if "r" in config else 1,
+                compute_linear_spec=config.model.lower() == "tacotron" or config.compute_linear_spec,
+                compute_f0=config.get("compute_f0", False),
+                f0_cache_path=config.get("f0_cache_path", None),
+                compute_energy=config.get("compute_energy", False),
+                energy_cache_path=config.get("energy_cache_path", None),
+                samples=samples,
+                ap=self.ap,
+                return_wav=config.return_wav if "return_wav" in config else False,
+                batch_group_size=0 if is_eval else config.batch_group_size * config.batch_size,
+                min_text_len=config.min_text_len,
+                max_text_len=config.max_text_len,
+                min_audio_len=config.min_audio_len,
+                max_audio_len=config.max_audio_len,
+                phoneme_cache_path=config.phoneme_cache_path,
+                precompute_num_workers=config.precompute_num_workers,
+                use_noise_augment=False if is_eval else config.use_noise_augment,
+                verbose=verbose,
+                speaker_id_mapping=speaker_id_mapping,
+                d_vector_mapping=d_vector_mapping if config.use_d_vector_file else None,
+                tokenizer=self.tokenizer,
+                start_by_longest=config.start_by_longest,
+                language_id_mapping=language_id_mapping,
+            )
+
+            # wait all the DDP process to be ready
+            if num_gpus > 1:
+                dist.barrier()
+
+            # sort input sequences from short to long
+            dataset.preprocess_samples()
+
+            # get samplers
+            sampler = self.get_sampler(config, dataset, num_gpus)
+
+            loader = DataLoader(
+                dataset,
+                batch_size=config.eval_batch_size if is_eval else config.batch_size,
+                shuffle=config.shuffle if sampler is None else False,  # if there is no other sampler
+                collate_fn=dataset.collate_fn,
+                drop_last=config.drop_last,  # setting this False might cause issues in AMP training.
+                sampler=sampler,
+                num_workers=config.num_eval_loader_workers if is_eval else config.num_loader_workers,
+                pin_memory=False,
+            )
+        return loader
+
+    def _get_test_aux_input(
+        self,
+    ) -> Dict:
+        d_vector = None
+        if self.config.use_d_vector_file:
+            d_vector = [self.speaker_manager.embeddings[name]["embedding"] for name in self.speaker_manager.embeddings]
+            d_vector = (random.sample(sorted(d_vector), 1),)
+
+        aux_inputs = {
+            "speaker_id": None
+            if not self.config.use_speaker_embedding
+            else random.sample(sorted(self.speaker_manager.name_to_id.values()), 1),
+            "d_vector": d_vector,
+            "style_wav": None,  # TODO: handle GST style input
+        }
+        return aux_inputs
+
+    def test_run(self, assets: Dict) -> Tuple[Dict, Dict]:
+        """Generic test run for `tts` models used by `Trainer`.
+
+        You can override this for a different behaviour.
+
+        Args:
+            assets (dict): A dict of training assets. For `tts` models, it must include `{'audio_processor': ap}`.
+
+        Returns:
+            Tuple[Dict, Dict]: Test figures and audios to be projected to Tensorboard.
+        """
+        print(" | > Synthesizing test sentences.")
+        test_audios = {}
+        test_figures = {}
+        test_sentences = self.config.test_sentences
+        aux_inputs = self._get_test_aux_input()
+        for idx, sen in enumerate(test_sentences):
+            if isinstance(sen, list):
+                aux_inputs = self.get_aux_input_from_test_sentences(sen)
+                sen = aux_inputs["text"]
+            outputs_dict = synthesis(
+                self,
+                sen,
+                self.config,
+                "cuda" in str(next(self.parameters()).device),
+                speaker_id=aux_inputs["speaker_id"],
+                d_vector=aux_inputs["d_vector"],
+                style_wav=aux_inputs["style_wav"],
+                use_griffin_lim=True,
+                do_trim_silence=False,
+            )
+            test_audios["{}-audio".format(idx)] = outputs_dict["wav"]
+            test_figures["{}-prediction".format(idx)] = plot_spectrogram(
+                outputs_dict["outputs"]["model_outputs"], self.ap, output_fig=False
+            )
+            test_figures["{}-alignment".format(idx)] = plot_alignment(
+                outputs_dict["outputs"]["alignments"], output_fig=False
+            )
+        return test_figures, test_audios
+
+    def on_init_start(self, trainer):
+        """Save the speaker.pth and language_ids.json at the beginning of the training. Also update both paths."""
+        if self.speaker_manager is not None:
+            output_path = os.path.join(trainer.output_path, "speakers.pth")
+            self.speaker_manager.save_ids_to_file(output_path)
+            trainer.config.speakers_file = output_path
+            # some models don't have `model_args` set
+            if hasattr(trainer.config, "model_args"):
+                trainer.config.model_args.speakers_file = output_path
+            trainer.config.save_json(os.path.join(trainer.output_path, "config.json"))
+            print(f" > `speakers.pth` is saved to {output_path}.")
+            print(" > `speakers_file` is updated in the config.json.")
+
+        if self.language_manager is not None:
+            output_path = os.path.join(trainer.output_path, "language_ids.json")
+            self.language_manager.save_ids_to_file(output_path)
+            trainer.config.language_ids_file = output_path
+            if hasattr(trainer.config, "model_args"):
+                trainer.config.model_args.language_ids_file = output_path
+            trainer.config.save_json(os.path.join(trainer.output_path, "config.json"))
+            print(f" > `language_ids.json` is saved to {output_path}.")
+            print(" > `language_ids_file` is updated in the config.json.")
+
+
+class BaseTTSE2E(BaseTTS):
+    def _set_model_args(self, config: Coqpit):
+        self.config = config
+        if "Config" in config.__class__.__name__:
+            num_chars = (
+                self.config.model_args.num_chars if self.tokenizer is None else self.tokenizer.characters.num_chars
+            )
+            self.config.model_args.num_chars = num_chars
+            self.config.num_chars = num_chars
+            self.args = config.model_args
+            self.args.num_chars = num_chars
+        elif "Args" in config.__class__.__name__:
+            self.args = config
+            self.args.num_chars = self.args.num_chars
+        else:
+            raise ValueError("config must be either a *Config or *Args")

TTS/tts/models/forward_tts.py

@@ -0,0 +1,862 @@
+from dataclasses import dataclass, field
+from typing import Dict, List, Tuple, Union
+
+import torch
+from coqpit import Coqpit
+from torch import nn
+from torch.cuda.amp.autocast_mode import autocast
+
+from TTS.tts.layers.feed_forward.decoder import Decoder
+from TTS.tts.layers.feed_forward.encoder import Encoder
+from TTS.tts.layers.generic.aligner import AlignmentNetwork
+from TTS.tts.layers.generic.pos_encoding import PositionalEncoding
+from TTS.tts.layers.glow_tts.duration_predictor import DurationPredictor
+from TTS.tts.models.base_tts import BaseTTS
+from TTS.tts.utils.helpers import average_over_durations, generate_path, maximum_path, sequence_mask
+from TTS.tts.utils.speakers import SpeakerManager
+from TTS.tts.utils.text.tokenizer import TTSTokenizer
+from TTS.tts.utils.visual import plot_alignment, plot_avg_energy, plot_avg_pitch, plot_spectrogram
+from TTS.utils.io import load_fsspec
+
+
+@dataclass
+class ForwardTTSArgs(Coqpit):
+    """ForwardTTS Model arguments.
+
+    Args:
+
+        num_chars (int):
+            Number of characters in the vocabulary. Defaults to 100.
+
+        out_channels (int):
+            Number of output channels. Defaults to 80.
+
+        hidden_channels (int):
+            Number of base hidden channels of the model. Defaults to 512.
+
+        use_aligner (bool):
+            Whether to use aligner network to learn the text to speech alignment or use pre-computed durations.
+            If set False, durations should be computed by `TTS/bin/compute_attention_masks.py` and path to the
+            pre-computed durations must be provided to `config.datasets[0].meta_file_attn_mask`. Defaults to True.
+
+        use_pitch (bool):
+            Use pitch predictor to learn the pitch. Defaults to True.
+
+        use_energy (bool):
+            Use energy predictor to learn the energy. Defaults to True.
+
+        duration_predictor_hidden_channels (int):
+            Number of hidden channels in the duration predictor. Defaults to 256.
+
+        duration_predictor_dropout_p (float):
+            Dropout rate for the duration predictor. Defaults to 0.1.
+
+        duration_predictor_kernel_size (int):
+            Kernel size of conv layers in the duration predictor. Defaults to 3.
+
+        pitch_predictor_hidden_channels (int):
+            Number of hidden channels in the pitch predictor. Defaults to 256.
+
+        pitch_predictor_dropout_p (float):
+            Dropout rate for the pitch predictor. Defaults to 0.1.
+
+        pitch_predictor_kernel_size (int):
+            Kernel size of conv layers in the pitch predictor. Defaults to 3.
+
+        pitch_embedding_kernel_size (int):
+            Kernel size of the projection layer in the pitch predictor. Defaults to 3.
+
+        energy_predictor_hidden_channels (int):
+            Number of hidden channels in the energy predictor. Defaults to 256.
+
+        energy_predictor_dropout_p (float):
+            Dropout rate for the energy predictor. Defaults to 0.1.
+
+        energy_predictor_kernel_size (int):
+            Kernel size of conv layers in the energy predictor. Defaults to 3.
+
+        energy_embedding_kernel_size (int):
+            Kernel size of the projection layer in the energy predictor. Defaults to 3.
+
+        positional_encoding (bool):
+            Whether to use positional encoding. Defaults to True.
+
+        positional_encoding_use_scale (bool):
+            Whether to use a learnable scale coeff in the positional encoding. Defaults to True.
+
+        length_scale (int):
+            Length scale that multiplies the predicted durations. Larger values result slower speech. Defaults to 1.0.
+
+        encoder_type (str):
+            Type of the encoder module. One of the encoders available in :class:`TTS.tts.layers.feed_forward.encoder`.
+            Defaults to `fftransformer` as in the paper.
+
+        encoder_params (dict):
+            Parameters of the encoder module. Defaults to ```{"hidden_channels_ffn": 1024, "num_heads": 1, "num_layers": 6, "dropout_p": 0.1}```
+
+        decoder_type (str):
+            Type of the decoder module. One of the decoders available in :class:`TTS.tts.layers.feed_forward.decoder`.
+            Defaults to `fftransformer` as in the paper.
+
+        decoder_params (str):
+            Parameters of the decoder module. Defaults to ```{"hidden_channels_ffn": 1024, "num_heads": 1, "num_layers": 6, "dropout_p": 0.1}```
+
+        detach_duration_predictor (bool):
+            Detach the input to the duration predictor from the earlier computation graph so that the duraiton loss
+            does not pass to the earlier layers. Defaults to True.
+
+        max_duration (int):
+            Maximum duration accepted by the model. Defaults to 75.
+
+        num_speakers (int):
+            Number of speakers for the speaker embedding layer. Defaults to 0.
+
+        speakers_file (str):
+            Path to the speaker mapping file for the Speaker Manager. Defaults to None.
+
+        speaker_embedding_channels (int):
+            Number of speaker embedding channels. Defaults to 256.
+
+        use_d_vector_file (bool):
+            Enable/Disable the use of d-vectors for multi-speaker training. Defaults to False.
+
+        d_vector_dim (int):
+            Number of d-vector channels. Defaults to 0.
+
+    """
+
+    num_chars: int = None
+    out_channels: int = 80
+    hidden_channels: int = 384
+    use_aligner: bool = True
+    # pitch params
+    use_pitch: bool = True
+    pitch_predictor_hidden_channels: int = 256
+    pitch_predictor_kernel_size: int = 3
+    pitch_predictor_dropout_p: float = 0.1
+    pitch_embedding_kernel_size: int = 3
+
+    # energy params
+    use_energy: bool = False
+    energy_predictor_hidden_channels: int = 256
+    energy_predictor_kernel_size: int = 3
+    energy_predictor_dropout_p: float = 0.1
+    energy_embedding_kernel_size: int = 3
+
+    # duration params
+    duration_predictor_hidden_channels: int = 256
+    duration_predictor_kernel_size: int = 3
+    duration_predictor_dropout_p: float = 0.1
+
+    positional_encoding: bool = True
+    poisitonal_encoding_use_scale: bool = True
+    length_scale: int = 1
+    encoder_type: str = "fftransformer"
+    encoder_params: dict = field(
+        default_factory=lambda: {"hidden_channels_ffn": 1024, "num_heads": 1, "num_layers": 6, "dropout_p": 0.1}
+    )
+    decoder_type: str = "fftransformer"
+    decoder_params: dict = field(
+        default_factory=lambda: {"hidden_channels_ffn": 1024, "num_heads": 1, "num_layers": 6, "dropout_p": 0.1}
+    )
+    detach_duration_predictor: bool = False
+    max_duration: int = 75
+    num_speakers: int = 1
+    use_speaker_embedding: bool = False
+    speakers_file: str = None
+    use_d_vector_file: bool = False
+    d_vector_dim: int = None
+    d_vector_file: str = None
+
+
+class ForwardTTS(BaseTTS):
+    """General forward TTS model implementation that uses an encoder-decoder architecture with an optional alignment
+    network and a pitch predictor.
+
+    If the alignment network is used, the model learns the text-to-speech alignment
+    from the data instead of using pre-computed durations.
+
+    If the pitch predictor is used, the model trains a pitch predictor that predicts average pitch value for each
+    input character as in the FastPitch model.
+
+    `ForwardTTS` can be configured to one of these architectures,
+
+        - FastPitch
+        - SpeedySpeech
+        - FastSpeech
+        - FastSpeech2 (requires average speech energy predictor)
+
+    Args:
+        config (Coqpit): Model coqpit class.
+        speaker_manager (SpeakerManager): Speaker manager for multi-speaker training. Only used for multi-speaker models.
+            Defaults to None.
+
+    Examples:
+        >>> from TTS.tts.models.fast_pitch import ForwardTTS, ForwardTTSArgs
+        >>> config = ForwardTTSArgs()
+        >>> model = ForwardTTS(config)
+    """
+
+    # pylint: disable=dangerous-default-value
+    def __init__(
+        self,
+        config: Coqpit,
+        ap: "AudioProcessor" = None,
+        tokenizer: "TTSTokenizer" = None,
+        speaker_manager: SpeakerManager = None,
+    ):
+        super().__init__(config, ap, tokenizer, speaker_manager)
+        self._set_model_args(config)
+
+        self.init_multispeaker(config)
+
+        self.max_duration = self.args.max_duration
+        self.use_aligner = self.args.use_aligner
+        self.use_pitch = self.args.use_pitch
+        self.use_energy = self.args.use_energy
+        self.binary_loss_weight = 0.0
+
+        self.length_scale = (
+            float(self.args.length_scale) if isinstance(self.args.length_scale, int) else self.args.length_scale
+        )
+
+        self.emb = nn.Embedding(self.args.num_chars, self.args.hidden_channels)
+
+        self.encoder = Encoder(
+            self.args.hidden_channels,
+            self.args.hidden_channels,
+            self.args.encoder_type,
+            self.args.encoder_params,
+            self.embedded_speaker_dim,
+        )
+
+        if self.args.positional_encoding:
+            self.pos_encoder = PositionalEncoding(self.args.hidden_channels)
+
+        self.decoder = Decoder(
+            self.args.out_channels,
+            self.args.hidden_channels,
+            self.args.decoder_type,
+            self.args.decoder_params,
+        )
+
+        self.duration_predictor = DurationPredictor(
+            self.args.hidden_channels,
+            self.args.duration_predictor_hidden_channels,
+            self.args.duration_predictor_kernel_size,
+            self.args.duration_predictor_dropout_p,
+        )
+
+        if self.args.use_pitch:
+            self.pitch_predictor = DurationPredictor(
+                self.args.hidden_channels,
+                self.args.pitch_predictor_hidden_channels,
+                self.args.pitch_predictor_kernel_size,
+                self.args.pitch_predictor_dropout_p,
+            )
+            self.pitch_emb = nn.Conv1d(
+                1,
+                self.args.hidden_channels,
+                kernel_size=self.args.pitch_embedding_kernel_size,
+                padding=int((self.args.pitch_embedding_kernel_size - 1) / 2),
+            )
+
+        if self.args.use_energy:
+            self.energy_predictor = DurationPredictor(
+                self.args.hidden_channels,
+                self.args.energy_predictor_hidden_channels,
+                self.args.energy_predictor_kernel_size,
+                self.args.energy_predictor_dropout_p,
+            )
+            self.energy_emb = nn.Conv1d(
+                1,
+                self.args.hidden_channels,
+                kernel_size=self.args.energy_embedding_kernel_size,
+                padding=int((self.args.energy_embedding_kernel_size - 1) / 2),
+            )
+
+        if self.args.use_aligner:
+            self.aligner = AlignmentNetwork(
+                in_query_channels=self.args.out_channels, in_key_channels=self.args.hidden_channels
+            )
+
+    def init_multispeaker(self, config: Coqpit):
+        """Init for multi-speaker training.
+
+        Args:
+            config (Coqpit): Model configuration.
+        """
+        self.embedded_speaker_dim = 0
+        # init speaker manager
+        if self.speaker_manager is None and (config.use_d_vector_file or config.use_speaker_embedding):
+            raise ValueError(
+                " > SpeakerManager is not provided. You must provide the SpeakerManager before initializing a multi-speaker model."
+            )
+        # set number of speakers
+        if self.speaker_manager is not None:
+            self.num_speakers = self.speaker_manager.num_speakers
+        # init d-vector embedding
+        if config.use_d_vector_file:
+            self.embedded_speaker_dim = config.d_vector_dim
+            if self.args.d_vector_dim != self.args.hidden_channels:
+                #self.proj_g = nn.Conv1d(self.args.d_vector_dim, self.args.hidden_channels, 1)
+                self.proj_g = nn.Linear(in_features=self.args.d_vector_dim, out_features=self.args.hidden_channels)
+        # init speaker embedding layer
+        if config.use_speaker_embedding and not config.use_d_vector_file:
+            print(" > Init speaker_embedding layer.")
+            self.emb_g = nn.Embedding(self.num_speakers, self.args.hidden_channels)
+            nn.init.uniform_(self.emb_g.weight, -0.1, 0.1)
+
+    @staticmethod
+    def generate_attn(dr, x_mask, y_mask=None):
+        """Generate an attention mask from the durations.
+
+        Shapes
+           - dr: :math:`(B, T_{en})`
+           - x_mask: :math:`(B, T_{en})`
+           - y_mask: :math:`(B, T_{de})`
+        """
+        # compute decode mask from the durations
+        if y_mask is None:
+            y_lengths = dr.sum(1).long()
+            y_lengths[y_lengths < 1] = 1
+            y_mask = torch.unsqueeze(sequence_mask(y_lengths, None), 1).to(dr.dtype)
+        attn_mask = torch.unsqueeze(x_mask, -1) * torch.unsqueeze(y_mask, 2)
+        attn = generate_path(dr, attn_mask.squeeze(1)).to(dr.dtype)
+        return attn
+
+    def expand_encoder_outputs(self, en, dr, x_mask, y_mask):
+        """Generate attention alignment map from durations and
+        expand encoder outputs
+
+        Shapes:
+            - en: :math:`(B, D_{en}, T_{en})`
+            - dr: :math:`(B, T_{en})`
+            - x_mask: :math:`(B, T_{en})`
+            - y_mask: :math:`(B, T_{de})`
+
+        Examples::
+
+            encoder output: [a,b,c,d]
+            durations: [1, 3, 2, 1]
+
+            expanded: [a, b, b, b, c, c, d]
+            attention map: [[0, 0, 0, 0, 0, 0, 1],
+                            [0, 0, 0, 0, 1, 1, 0],
+                            [0, 1, 1, 1, 0, 0, 0],
+                            [1, 0, 0, 0, 0, 0, 0]]
+        """
+        attn = self.generate_attn(dr, x_mask, y_mask)
+        o_en_ex = torch.matmul(attn.squeeze(1).transpose(1, 2).to(en.dtype), en.transpose(1, 2)).transpose(1, 2)
+        return o_en_ex, attn
+
+    def format_durations(self, o_dr_log, x_mask):
+        """Format predicted durations.
+        1. Convert to linear scale from log scale
+        2. Apply the length scale for speed adjustment
+        3. Apply masking.
+        4. Cast 0 durations to 1.
+        5. Round the duration values.
+
+        Args:
+            o_dr_log: Log scale durations.
+            x_mask: Input text mask.
+
+        Shapes:
+            - o_dr_log: :math:`(B, T_{de})`
+            - x_mask: :math:`(B, T_{en})`
+        """
+        o_dr = (torch.exp(o_dr_log) - 1) * x_mask * self.length_scale
+        o_dr[o_dr < 1] = 1.0
+        o_dr = torch.round(o_dr)
+        return o_dr
+
+    def _forward_encoder(
+        self, x: torch.LongTensor, x_mask: torch.FloatTensor, g: torch.FloatTensor = None
+    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
+        """Encoding forward pass.
+
+        1. Embed speaker IDs if multi-speaker mode.
+        2. Embed character sequences.
+        3. Run the encoder network.
+        4. Sum encoder outputs and speaker embeddings
+
+        Args:
+            x (torch.LongTensor): Input sequence IDs.
+            x_mask (torch.FloatTensor): Input squence mask.
+            g (torch.FloatTensor, optional): Conditioning vectors. In general speaker embeddings. Defaults to None.
+
+        Returns:
+            Tuple[torch.tensor, torch.tensor, torch.tensor, torch.tensor, torch.tensor]:
+                encoder output, encoder output for the duration predictor, input sequence mask, speaker embeddings,
+                character embeddings
+
+        Shapes:
+            - x: :math:`(B, T_{en})`
+            - x_mask: :math:`(B, 1, T_{en})`
+            - g: :math:`(B, C)`
+        """
+        if hasattr(self, "emb_g"):
+            g = g.type(torch.LongTensor)
+            g = self.emb_g(g)  # [B, C, 1]
+        if g is not None:
+            g = g.unsqueeze(-1)
+        # [B, T, C]
+        x_emb = self.emb(x)
+        # encoder pass
+	#o_en = self.encoder(torch.transpose(x_emb, 1, -1), x_mask)
+        o_en = self.encoder(torch.transpose(x_emb, 1, -1), x_mask, g)
+        # speaker conditioning
+        # TODO: try different ways of conditioning
+        if g is not None: 
+            if hasattr(self, "proj_g"):
+                g = self.proj_g(g.view(g.shape[0], -1)).unsqueeze(-1)            
+            o_en = o_en + g
+        return o_en, x_mask, g, x_emb
+
+    def _forward_decoder(
+        self,
+        o_en: torch.FloatTensor,
+        dr: torch.IntTensor,
+        x_mask: torch.FloatTensor,
+        y_lengths: torch.IntTensor,
+        g: torch.FloatTensor,
+    ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+        """Decoding forward pass.
+
+        1. Compute the decoder output mask
+        2. Expand encoder output with the durations.
+        3. Apply position encoding.
+        4. Add speaker embeddings if multi-speaker mode.
+        5. Run the decoder.
+
+        Args:
+            o_en (torch.FloatTensor): Encoder output.
+            dr (torch.IntTensor): Ground truth durations or alignment network durations.
+            x_mask (torch.IntTensor): Input sequence mask.
+            y_lengths (torch.IntTensor): Output sequence lengths.
+            g (torch.FloatTensor): Conditioning vectors. In general speaker embeddings.
+
+        Returns:
+            Tuple[torch.FloatTensor, torch.FloatTensor]: Decoder output, attention map from durations.
+        """
+        y_mask = torch.unsqueeze(sequence_mask(y_lengths, None), 1).to(o_en.dtype)
+        # expand o_en with durations
+        o_en_ex, attn = self.expand_encoder_outputs(o_en, dr, x_mask, y_mask)
+        # positional encoding
+        if hasattr(self, "pos_encoder"):
+            o_en_ex = self.pos_encoder(o_en_ex, y_mask)
+        # decoder pass
+        o_de = self.decoder(o_en_ex, y_mask, g=g)
+        return o_de.transpose(1, 2), attn.transpose(1, 2)
+
+    def _forward_pitch_predictor(
+        self,
+        o_en: torch.FloatTensor,
+        x_mask: torch.IntTensor,
+        pitch: torch.FloatTensor = None,
+        dr: torch.IntTensor = None,
+    ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+        """Pitch predictor forward pass.
+
+        1. Predict pitch from encoder outputs.
+        2. In training - Compute average pitch values for each input character from the ground truth pitch values.
+        3. Embed average pitch values.
+
+        Args:
+            o_en (torch.FloatTensor): Encoder output.
+            x_mask (torch.IntTensor): Input sequence mask.
+            pitch (torch.FloatTensor, optional): Ground truth pitch values. Defaults to None.
+            dr (torch.IntTensor, optional): Ground truth durations. Defaults to None.
+
+        Returns:
+            Tuple[torch.FloatTensor, torch.FloatTensor]: Pitch embedding, pitch prediction.
+
+        Shapes:
+            - o_en: :math:`(B, C, T_{en})`
+            - x_mask: :math:`(B, 1, T_{en})`
+            - pitch: :math:`(B, 1, T_{de})`
+            - dr: :math:`(B, T_{en})`
+        """
+        o_pitch = self.pitch_predictor(o_en, x_mask)
+        if pitch is not None:
+            avg_pitch = average_over_durations(pitch, dr)
+            o_pitch_emb = self.pitch_emb(avg_pitch)
+            return o_pitch_emb, o_pitch, avg_pitch
+        o_pitch_emb = self.pitch_emb(o_pitch)
+        return o_pitch_emb, o_pitch
+
+    def _forward_energy_predictor(
+        self,
+        o_en: torch.FloatTensor,
+        x_mask: torch.IntTensor,
+        energy: torch.FloatTensor = None,
+        dr: torch.IntTensor = None,
+    ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+        """Energy predictor forward pass.
+
+        1. Predict energy from encoder outputs.
+        2. In training - Compute average pitch values for each input character from the ground truth pitch values.
+        3. Embed average energy values.
+
+        Args:
+            o_en (torch.FloatTensor): Encoder output.
+            x_mask (torch.IntTensor): Input sequence mask.
+            energy (torch.FloatTensor, optional): Ground truth energy values. Defaults to None.
+            dr (torch.IntTensor, optional): Ground truth durations. Defaults to None.
+
+        Returns:
+            Tuple[torch.FloatTensor, torch.FloatTensor]: Energy embedding, energy prediction.
+
+        Shapes:
+            - o_en: :math:`(B, C, T_{en})`
+            - x_mask: :math:`(B, 1, T_{en})`
+            - pitch: :math:`(B, 1, T_{de})`
+            - dr: :math:`(B, T_{en})`
+        """
+        o_energy = self.energy_predictor(o_en, x_mask)
+        if energy is not None:
+            avg_energy = average_over_durations(energy, dr)
+            o_energy_emb = self.energy_emb(avg_energy)
+            return o_energy_emb, o_energy, avg_energy
+        o_energy_emb = self.energy_emb(o_energy)
+        return o_energy_emb, o_energy
+
+    def _forward_aligner(
+        self, x: torch.FloatTensor, y: torch.FloatTensor, x_mask: torch.IntTensor, y_mask: torch.IntTensor
+    ) -> Tuple[torch.IntTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
+        """Aligner forward pass.
+
+        1. Compute a mask to apply to the attention map.
+        2. Run the alignment network.
+        3. Apply MAS to compute the hard alignment map.
+        4. Compute the durations from the hard alignment map.
+
+        Args:
+            x (torch.FloatTensor): Input sequence.
+            y (torch.FloatTensor): Output sequence.
+            x_mask (torch.IntTensor): Input sequence mask.
+            y_mask (torch.IntTensor): Output sequence mask.
+
+        Returns:
+            Tuple[torch.IntTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
+                Durations from the hard alignment map, soft alignment potentials, log scale alignment potentials,
+                hard alignment map.
+
+        Shapes:
+            - x: :math:`[B, T_en, C_en]`
+            - y: :math:`[B, T_de, C_de]`
+            - x_mask: :math:`[B, 1, T_en]`
+            - y_mask: :math:`[B, 1, T_de]`
+
+            - o_alignment_dur: :math:`[B, T_en]`
+            - alignment_soft: :math:`[B, T_en, T_de]`
+            - alignment_logprob: :math:`[B, 1, T_de, T_en]`
+            - alignment_mas: :math:`[B, T_en, T_de]`
+        """
+        attn_mask = torch.unsqueeze(x_mask, -1) * torch.unsqueeze(y_mask, 2)
+        alignment_soft, alignment_logprob = self.aligner(y.transpose(1, 2), x.transpose(1, 2), x_mask, None)
+        alignment_mas = maximum_path(
+            alignment_soft.squeeze(1).transpose(1, 2).contiguous(), attn_mask.squeeze(1).contiguous()
+        )
+        o_alignment_dur = torch.sum(alignment_mas, -1).int()
+        alignment_soft = alignment_soft.squeeze(1).transpose(1, 2)
+        return o_alignment_dur, alignment_soft, alignment_logprob, alignment_mas
+
+    def _set_speaker_input(self, aux_input: Dict):
+        d_vectors = aux_input.get("d_vectors", None)
+        speaker_ids = aux_input.get("speaker_ids", None)
+
+        if d_vectors is not None and speaker_ids is not None:
+            raise ValueError("[!] Cannot use d-vectors and speaker-ids together.")
+
+        if speaker_ids is not None and not hasattr(self, "emb_g"):
+            raise ValueError("[!] Cannot use speaker-ids without enabling speaker embedding.")
+
+        g = speaker_ids if speaker_ids is not None else d_vectors
+        return g
+
+    def forward(
+        self,
+        x: torch.LongTensor,
+        x_lengths: torch.LongTensor,
+        y_lengths: torch.LongTensor,
+        y: torch.FloatTensor = None,
+        dr: torch.IntTensor = None,
+        pitch: torch.FloatTensor = None,
+        energy: torch.FloatTensor = None,
+        aux_input: Dict = {"d_vectors": None, "speaker_ids": None},  # pylint: disable=unused-argument
+    ) -> Dict:
+        """Model's forward pass.
+
+        Args:
+            x (torch.LongTensor): Input character sequences.
+            x_lengths (torch.LongTensor): Input sequence lengths.
+            y_lengths (torch.LongTensor): Output sequnce lengths. Defaults to None.
+            y (torch.FloatTensor): Spectrogram frames. Only used when the alignment network is on. Defaults to None.
+            dr (torch.IntTensor): Character durations over the spectrogram frames. Only used when the alignment network is off. Defaults to None.
+            pitch (torch.FloatTensor): Pitch values for each spectrogram frame. Only used when the pitch predictor is on. Defaults to None.
+            energy (torch.FloatTensor): energy values for each spectrogram frame. Only used when the energy predictor is on. Defaults to None.
+            aux_input (Dict): Auxiliary model inputs for multi-speaker training. Defaults to `{"d_vectors": 0, "speaker_ids": None}`.
+
+        Shapes:
+            - x: :math:`[B, T_max]`
+            - x_lengths: :math:`[B]`
+            - y_lengths: :math:`[B]`
+            - y: :math:`[B, T_max2]`
+            - dr: :math:`[B, T_max]`
+            - g: :math:`[B, C]`
+            - pitch: :math:`[B, 1, T]`
+        """
+        g = self._set_speaker_input(aux_input)
+        # compute sequence masks
+        y_mask = torch.unsqueeze(sequence_mask(y_lengths, None), 1).float()
+        x_mask = torch.unsqueeze(sequence_mask(x_lengths, x.shape[1]), 1).float()
+        # encoder pass
+        o_en, x_mask, g, x_emb = self._forward_encoder(x, x_mask, g)
+        # duration predictor pass
+        if self.args.detach_duration_predictor:
+            o_dr_log = self.duration_predictor(o_en.detach(), x_mask)
+        else:
+            o_dr_log = self.duration_predictor(o_en, x_mask)
+        o_dr = torch.clamp(torch.exp(o_dr_log) - 1, 0, self.max_duration)
+        # generate attn mask from predicted durations
+        o_attn = self.generate_attn(o_dr.squeeze(1), x_mask)
+        # aligner
+        o_alignment_dur = None
+        alignment_soft = None
+        alignment_logprob = None
+        alignment_mas = None
+        if self.use_aligner:
+            o_alignment_dur, alignment_soft, alignment_logprob, alignment_mas = self._forward_aligner(
+                x_emb, y, x_mask, y_mask
+            )
+            alignment_soft = alignment_soft.transpose(1, 2)
+            alignment_mas = alignment_mas.transpose(1, 2)
+            dr = o_alignment_dur
+        # pitch predictor pass
+        o_pitch = None
+        avg_pitch = None
+        if self.args.use_pitch:
+            o_pitch_emb, o_pitch, avg_pitch = self._forward_pitch_predictor(o_en, x_mask, pitch, dr)
+            o_en = o_en + o_pitch_emb
+        # energy predictor pass
+        o_energy = None
+        avg_energy = None
+        if self.args.use_energy:
+            o_energy_emb, o_energy, avg_energy = self._forward_energy_predictor(o_en, x_mask, energy, dr)
+            o_en = o_en + o_energy_emb
+        # decoder pass
+        o_de, attn = self._forward_decoder(
+            o_en, dr, x_mask, y_lengths, g=None
+        )  # TODO: maybe pass speaker embedding (g) too
+        outputs = {
+            "model_outputs": o_de,  # [B, T, C]
+            "durations_log": o_dr_log.squeeze(1),  # [B, T]
+            "durations": o_dr.squeeze(1),  # [B, T]
+            "attn_durations": o_attn,  # for visualization [B, T_en, T_de']
+            "pitch_avg": o_pitch,
+            "pitch_avg_gt": avg_pitch,
+            "energy_avg": o_energy,
+            "energy_avg_gt": avg_energy,
+            "alignments": attn,  # [B, T_de, T_en]
+            "alignment_soft": alignment_soft,
+            "alignment_mas": alignment_mas,
+            "o_alignment_dur": o_alignment_dur,
+            "alignment_logprob": alignment_logprob,
+            "x_mask": x_mask,
+            "y_mask": y_mask,
+        }
+        return outputs
+
+    @torch.no_grad()
+    def inference(self, x, aux_input={"d_vectors": None, "speaker_ids": None}):  # pylint: disable=unused-argument
+        """Model's inference pass.
+
+        Args:
+            x (torch.LongTensor): Input character sequence.
+            aux_input (Dict): Auxiliary model inputs. Defaults to `{"d_vectors": None, "speaker_ids": None}`.
+
+        Shapes:
+            - x: [B, T_max]
+            - x_lengths: [B]
+            - g: [B, C]
+        """
+        g = self._set_speaker_input(aux_input)
+        x_lengths = torch.tensor(x.shape[1:2]).to(x.device)
+        x_mask = torch.unsqueeze(sequence_mask(x_lengths, x.shape[1]), 1).to(x.dtype).float()
+        # encoder pass
+        o_en, x_mask, g, _ = self._forward_encoder(x, x_mask, g)
+        # duration predictor pass
+        o_dr_log = self.duration_predictor(o_en.squeeze(), x_mask)
+        o_dr = self.format_durations(o_dr_log, x_mask).squeeze(1)
+        y_lengths = o_dr.sum(1)
+
+        # pitch predictor pass
+        o_pitch = None
+        if self.args.use_pitch:
+            o_pitch_emb, o_pitch = self._forward_pitch_predictor(o_en, x_mask)
+            o_en = o_en + o_pitch_emb
+        # energy predictor pass
+        o_energy = None
+        if self.args.use_energy:
+            o_energy_emb, o_energy = self._forward_energy_predictor(o_en, x_mask)
+            o_en = o_en + o_energy_emb
+        # decoder pass
+        o_de, attn = self._forward_decoder(o_en, o_dr, x_mask, y_lengths, g=None)
+        outputs = {
+            "model_outputs": o_de,
+            "alignments": attn,
+            "pitch": o_pitch,
+            "energy": o_energy,
+            "durations_log": o_dr_log,
+        }
+        return outputs
+
+    def train_step(self, batch: dict, criterion: nn.Module):
+        text_input = batch["text_input"]
+        text_lengths = batch["text_lengths"]
+        mel_input = batch["mel_input"]
+        mel_lengths = batch["mel_lengths"]
+        pitch = batch["pitch"] if self.args.use_pitch else None
+        energy = batch["energy"] if self.args.use_energy else None
+        d_vectors = batch["d_vectors"]
+        speaker_ids = batch["speaker_ids"]
+        durations = batch["durations"]
+        aux_input = {"d_vectors": d_vectors, "speaker_ids": speaker_ids}
+
+        # forward pass
+        outputs = self.forward(
+            text_input,
+            text_lengths,
+            mel_lengths,
+            y=mel_input,
+            dr=durations,
+            pitch=pitch,
+            energy=energy,
+            aux_input=aux_input,
+        )
+        # use aligner's output as the duration target
+        if self.use_aligner:
+            durations = outputs["o_alignment_dur"]
+        # use float32 in AMP
+        with autocast(enabled=False):
+            # compute loss
+            loss_dict = criterion(
+                decoder_output=outputs["model_outputs"],
+                decoder_target=mel_input,
+                decoder_output_lens=mel_lengths,
+                dur_output=outputs["durations_log"],
+                dur_target=durations,
+                pitch_output=outputs["pitch_avg"] if self.use_pitch else None,
+                pitch_target=outputs["pitch_avg_gt"] if self.use_pitch else None,
+                energy_output=outputs["energy_avg"] if self.use_energy else None,
+                energy_target=outputs["energy_avg_gt"] if self.use_energy else None,
+                input_lens=text_lengths,
+                alignment_logprob=outputs["alignment_logprob"] if self.use_aligner else None,
+                alignment_soft=outputs["alignment_soft"],
+                alignment_hard=outputs["alignment_mas"],
+                binary_loss_weight=self.binary_loss_weight,
+            )
+            # compute duration error
+            durations_pred = outputs["durations"]
+            duration_error = torch.abs(durations - durations_pred).sum() / text_lengths.sum()
+            loss_dict["duration_error"] = duration_error
+
+        return outputs, loss_dict
+
+    def _create_logs(self, batch, outputs, ap):
+        """Create common logger outputs."""
+        model_outputs = outputs["model_outputs"]
+        alignments = outputs["alignments"]
+        mel_input = batch["mel_input"]
+
+        pred_spec = model_outputs[0].data.cpu().numpy()
+        gt_spec = mel_input[0].data.cpu().numpy()
+        align_img = alignments[0].data.cpu().numpy()
+
+        figures = {
+            "prediction": plot_spectrogram(pred_spec, ap, output_fig=False),
+            "ground_truth": plot_spectrogram(gt_spec, ap, output_fig=False),
+            "alignment": plot_alignment(align_img, output_fig=False),
+        }
+
+        # plot pitch figures
+        if self.args.use_pitch:
+            pitch_avg = abs(outputs["pitch_avg_gt"][0, 0].data.cpu().numpy())
+            pitch_avg_hat = abs(outputs["pitch_avg"][0, 0].data.cpu().numpy())
+            chars = self.tokenizer.decode(batch["text_input"][0].data.cpu().numpy())
+            pitch_figures = {
+                "pitch_ground_truth": plot_avg_pitch(pitch_avg, chars, output_fig=False),
+                "pitch_avg_predicted": plot_avg_pitch(pitch_avg_hat, chars, output_fig=False),
+            }
+            figures.update(pitch_figures)
+
+        # plot energy figures
+        if self.args.use_energy:
+            energy_avg = abs(outputs["energy_avg_gt"][0, 0].data.cpu().numpy())
+            energy_avg_hat = abs(outputs["energy_avg"][0, 0].data.cpu().numpy())
+            chars = self.tokenizer.decode(batch["text_input"][0].data.cpu().numpy())
+            energy_figures = {
+                "energy_ground_truth": plot_avg_energy(energy_avg, chars, output_fig=False),
+                "energy_avg_predicted": plot_avg_energy(energy_avg_hat, chars, output_fig=False),
+            }
+            figures.update(energy_figures)
+
+        # plot the attention mask computed from the predicted durations
+        if "attn_durations" in outputs:
+            alignments_hat = outputs["attn_durations"][0].data.cpu().numpy()
+            figures["alignment_hat"] = plot_alignment(alignments_hat.T, output_fig=False)
+
+        # Sample audio
+        train_audio = ap.inv_melspectrogram(pred_spec.T)
+        return figures, {"audio": train_audio}
+
+    def train_log(
+        self, batch: dict, outputs: dict, logger: "Logger", assets: dict, steps: int
+    ) -> None:  # pylint: disable=no-self-use
+        figures, audios = self._create_logs(batch, outputs, self.ap)
+        logger.train_figures(steps, figures)
+        logger.train_audios(steps, audios, self.ap.sample_rate)
+
+    def eval_step(self, batch: dict, criterion: nn.Module):
+        return self.train_step(batch, criterion)
+
+    def eval_log(self, batch: dict, outputs: dict, logger: "Logger", assets: dict, steps: int) -> None:
+        figures, audios = self._create_logs(batch, outputs, self.ap)
+        logger.eval_figures(steps, figures)
+        logger.eval_audios(steps, audios, self.ap.sample_rate)
+
+    def load_checkpoint(
+        self, config, checkpoint_path, eval=False, cache=False
+    ):  # pylint: disable=unused-argument, redefined-builtin
+        state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"), cache=cache)
+        self.load_state_dict(state["model"])
+        if eval:
+            self.eval()
+            assert not self.training
+
+    def get_criterion(self):
+        from TTS.tts.layers.losses import ForwardTTSLoss  # pylint: disable=import-outside-toplevel
+
+        return ForwardTTSLoss(self.config)
+
+    def on_train_step_start(self, trainer):
+        """Schedule binary loss weight."""
+        self.binary_loss_weight = min(trainer.epochs_done / self.config.binary_loss_warmup_epochs, 1.0) * 1.0
+
+    @staticmethod
+    def init_from_config(config: "ForwardTTSConfig", samples: Union[List[List], List[Dict]] = None):
+        """Initiate model from config
+
+        Args:
+            config (ForwardTTSConfig): Model config.
+            samples (Union[List[List], List[Dict]]): Training samples to parse speaker ids for training.
+                Defaults to None.
+        """
+        from TTS.utils.audio import AudioProcessor
+
+        ap = AudioProcessor.init_from_config(config)
+        tokenizer, new_config = TTSTokenizer.init_from_config(config)
+        speaker_manager = SpeakerManager.init_from_config(config, samples)
+        return ForwardTTS(new_config, ap, tokenizer, speaker_manager)

TTS/tts/models/glow_tts.py

@@ -0,0 +1,557 @@
+import math
+from typing import Dict, List, Tuple, Union
+
+import torch
+from coqpit import Coqpit
+from torch import nn
+from torch.cuda.amp.autocast_mode import autocast
+from torch.nn import functional as F
+
+from TTS.tts.configs.glow_tts_config import GlowTTSConfig
+from TTS.tts.layers.glow_tts.decoder import Decoder
+from TTS.tts.layers.glow_tts.encoder import Encoder
+from TTS.tts.models.base_tts import BaseTTS
+from TTS.tts.utils.helpers import generate_path, maximum_path, sequence_mask
+from TTS.tts.utils.speakers import SpeakerManager
+from TTS.tts.utils.synthesis import synthesis
+from TTS.tts.utils.text.tokenizer import TTSTokenizer
+from TTS.tts.utils.visual import plot_alignment, plot_spectrogram
+from TTS.utils.io import load_fsspec
+
+
+class GlowTTS(BaseTTS):
+    """GlowTTS model.
+
+    Paper::
+        https://arxiv.org/abs/2005.11129
+
+    Paper abstract::
+        Recently, text-to-speech (TTS) models such as FastSpeech and ParaNet have been proposed to generate
+        mel-spectrograms from text in parallel. Despite the advantage, the parallel TTS models cannot be trained
+        without guidance from autoregressive TTS models as their external aligners. In this work, we propose Glow-TTS,
+        a flow-based generative model for parallel TTS that does not require any external aligner. By combining the
+        properties of flows and dynamic programming, the proposed model searches for the most probable monotonic
+        alignment between text and the latent representation of speech on its own. We demonstrate that enforcing hard
+        monotonic alignments enables robust TTS, which generalizes to long utterances, and employing generative flows
+        enables fast, diverse, and controllable speech synthesis. Glow-TTS obtains an order-of-magnitude speed-up over
+        the autoregressive model, Tacotron 2, at synthesis with comparable speech quality. We further show that our
+        model can be easily extended to a multi-speaker setting.
+
+    Check :class:`TTS.tts.configs.glow_tts_config.GlowTTSConfig` for class arguments.
+
+    Examples:
+        Init only model layers.
+
+        >>> from TTS.tts.configs.glow_tts_config import GlowTTSConfig
+        >>> from TTS.tts.models.glow_tts import GlowTTS
+        >>> config = GlowTTSConfig(num_chars=2)
+        >>> model = GlowTTS(config)
+
+        Fully init a model ready for action. All the class attributes and class members
+        (e.g Tokenizer, AudioProcessor, etc.). are initialized internally based on config values.
+
+        >>> from TTS.tts.configs.glow_tts_config import GlowTTSConfig
+        >>> from TTS.tts.models.glow_tts import GlowTTS
+        >>> config = GlowTTSConfig()
+        >>> model = GlowTTS.init_from_config(config, verbose=False)
+    """
+
+    def __init__(
+        self,
+        config: GlowTTSConfig,
+        ap: "AudioProcessor" = None,
+        tokenizer: "TTSTokenizer" = None,
+        speaker_manager: SpeakerManager = None,
+    ):
+        super().__init__(config, ap, tokenizer, speaker_manager)
+
+        # pass all config fields to `self`
+        # for fewer code change
+        self.config = config
+        for key in config:
+            setattr(self, key, config[key])
+
+        self.decoder_output_dim = config.out_channels
+
+        # init multi-speaker layers if necessary
+        self.init_multispeaker(config)
+
+        self.run_data_dep_init = config.data_dep_init_steps > 0
+        self.encoder = Encoder(
+            self.num_chars,
+            out_channels=self.out_channels,
+            hidden_channels=self.hidden_channels_enc,
+            hidden_channels_dp=self.hidden_channels_dp,
+            encoder_type=self.encoder_type,
+            encoder_params=self.encoder_params,
+            mean_only=self.mean_only,
+            use_prenet=self.use_encoder_prenet,
+            dropout_p_dp=self.dropout_p_dp,
+            c_in_channels=self.c_in_channels,
+        )
+
+        self.decoder = Decoder(
+            self.out_channels,
+            self.hidden_channels_dec,
+            self.kernel_size_dec,
+            self.dilation_rate,
+            self.num_flow_blocks_dec,
+            self.num_block_layers,
+            dropout_p=self.dropout_p_dec,
+            num_splits=self.num_splits,
+            num_squeeze=self.num_squeeze,
+            sigmoid_scale=self.sigmoid_scale,
+            c_in_channels=self.c_in_channels,
+        )
+
+    def init_multispeaker(self, config: Coqpit):
+        """Init speaker embedding layer if `use_speaker_embedding` is True and set the expected speaker embedding
+        vector dimension to the encoder layer channel size. If model uses d-vectors, then it only sets
+        speaker embedding vector dimension to the d-vector dimension from the config.
+
+        Args:
+            config (Coqpit): Model configuration.
+        """
+        self.embedded_speaker_dim = 0
+        # set number of speakers - if num_speakers is set in config, use it, otherwise use speaker_manager
+        if self.speaker_manager is not None:
+            self.num_speakers = self.speaker_manager.num_speakers
+        # set ultimate speaker embedding size
+        if config.use_d_vector_file:
+            self.embedded_speaker_dim = (
+                config.d_vector_dim if "d_vector_dim" in config and config.d_vector_dim is not None else 512
+            )
+            if self.speaker_manager is not None:
+                assert (
+                    config.d_vector_dim == self.speaker_manager.embedding_dim
+                ), " [!] d-vector dimension mismatch b/w config and speaker manager."
+        # init speaker embedding layer
+        if config.use_speaker_embedding and not config.use_d_vector_file:
+            print(" > Init speaker_embedding layer.")
+            self.embedded_speaker_dim = self.hidden_channels_enc
+            self.emb_g = nn.Embedding(self.num_speakers, self.hidden_channels_enc)
+            nn.init.uniform_(self.emb_g.weight, -0.1, 0.1)
+        # set conditioning dimensions
+        self.c_in_channels = self.embedded_speaker_dim
+
+    @staticmethod
+    def compute_outputs(attn, o_mean, o_log_scale, x_mask):
+        """Compute and format the mode outputs with the given alignment map"""
+        y_mean = torch.matmul(attn.squeeze(1).transpose(1, 2), o_mean.transpose(1, 2)).transpose(
+            1, 2
+        )  # [b, t', t], [b, t, d] -> [b, d, t']
+        y_log_scale = torch.matmul(attn.squeeze(1).transpose(1, 2), o_log_scale.transpose(1, 2)).transpose(
+            1, 2
+        )  # [b, t', t], [b, t, d] -> [b, d, t']
+        # compute total duration with adjustment
+        o_attn_dur = torch.log(1 + torch.sum(attn, -1)) * x_mask
+        return y_mean, y_log_scale, o_attn_dur
+
+    def unlock_act_norm_layers(self):
+        """Unlock activation normalization layers for data depended initalization."""
+        for f in self.decoder.flows:
+            if getattr(f, "set_ddi", False):
+                f.set_ddi(True)
+
+    def lock_act_norm_layers(self):
+        """Lock activation normalization layers."""
+        for f in self.decoder.flows:
+            if getattr(f, "set_ddi", False):
+                f.set_ddi(False)
+
+    def _set_speaker_input(self, aux_input: Dict):
+        if aux_input is None:
+            d_vectors = None
+            speaker_ids = None
+        else:
+            d_vectors = aux_input.get("d_vectors", None)
+            speaker_ids = aux_input.get("speaker_ids", None)
+
+        if d_vectors is not None and speaker_ids is not None:
+            raise ValueError("[!] Cannot use d-vectors and speaker-ids together.")
+
+        if speaker_ids is not None and not hasattr(self, "emb_g"):
+            raise ValueError("[!] Cannot use speaker-ids without enabling speaker embedding.")
+
+        g = speaker_ids if speaker_ids is not None else d_vectors
+        return g
+
+    def _speaker_embedding(self, aux_input: Dict) -> Union[torch.tensor, None]:
+        g = self._set_speaker_input(aux_input)
+        # speaker embedding
+        if g is not None:
+            if hasattr(self, "emb_g"):
+                # use speaker embedding layer
+                if not g.size():  # if is a scalar
+                    g = g.unsqueeze(0)  # unsqueeze
+                g = F.normalize(self.emb_g(g)).unsqueeze(-1)  # [b, h, 1]
+            else:
+                # use d-vector
+                g = F.normalize(g).unsqueeze(-1)  # [b, h, 1]
+        return g
+
+    def forward(
+        self, x, x_lengths, y, y_lengths=None, aux_input={"d_vectors": None, "speaker_ids": None}
+    ):  # pylint: disable=dangerous-default-value
+        """
+        Args:
+            x (torch.Tensor):
+                Input text sequence ids. :math:`[B, T_en]`
+
+            x_lengths (torch.Tensor):
+                Lengths of input text sequences. :math:`[B]`
+
+            y (torch.Tensor):
+                Target mel-spectrogram frames. :math:`[B, T_de, C_mel]`
+
+            y_lengths (torch.Tensor):
+                Lengths of target mel-spectrogram frames. :math:`[B]`
+
+            aux_input (Dict):
+                Auxiliary inputs. `d_vectors` is speaker embedding vectors for a multi-speaker model.
+                :math:`[B, D_vec]`. `speaker_ids` is speaker ids for a multi-speaker model usind speaker-embedding
+                layer. :math:`B`
+
+        Returns:
+            Dict:
+                - z: :math: `[B, T_de, C]`
+                - logdet: :math:`B`
+                - y_mean: :math:`[B, T_de, C]`
+                - y_log_scale: :math:`[B, T_de, C]`
+                - alignments: :math:`[B, T_en, T_de]`
+                - durations_log: :math:`[B, T_en, 1]`
+                - total_durations_log: :math:`[B, T_en, 1]`
+        """
+        # [B, T, C] -> [B, C, T]
+        y = y.transpose(1, 2)
+        y_max_length = y.size(2)
+        # norm speaker embeddings
+        g = self._speaker_embedding(aux_input)
+        # embedding pass
+        o_mean, o_log_scale, o_dur_log, x_mask = self.encoder(x, x_lengths, g=g)
+        # drop redisual frames wrt num_squeeze and set y_lengths.
+        y, y_lengths, y_max_length, attn = self.preprocess(y, y_lengths, y_max_length, None)
+        # create masks
+        y_mask = torch.unsqueeze(sequence_mask(y_lengths, y_max_length), 1).to(x_mask.dtype)
+        # [B, 1, T_en, T_de]
+        attn_mask = torch.unsqueeze(x_mask, -1) * torch.unsqueeze(y_mask, 2)
+        # decoder pass
+        z, logdet = self.decoder(y, y_mask, g=g, reverse=False)
+        # find the alignment path
+        with torch.no_grad():
+            o_scale = torch.exp(-2 * o_log_scale)
+            logp1 = torch.sum(-0.5 * math.log(2 * math.pi) - o_log_scale, [1]).unsqueeze(-1)  # [b, t, 1]
+            logp2 = torch.matmul(o_scale.transpose(1, 2), -0.5 * (z**2))  # [b, t, d] x [b, d, t'] = [b, t, t']
+            logp3 = torch.matmul((o_mean * o_scale).transpose(1, 2), z)  # [b, t, d] x [b, d, t'] = [b, t, t']
+            logp4 = torch.sum(-0.5 * (o_mean**2) * o_scale, [1]).unsqueeze(-1)  # [b, t, 1]
+            logp = logp1 + logp2 + logp3 + logp4  # [b, t, t']
+            attn = maximum_path(logp, attn_mask.squeeze(1)).unsqueeze(1).detach()
+        y_mean, y_log_scale, o_attn_dur = self.compute_outputs(attn, o_mean, o_log_scale, x_mask)
+        attn = attn.squeeze(1).permute(0, 2, 1)
+        outputs = {
+            "z": z.transpose(1, 2),
+            "logdet": logdet,
+            "y_mean": y_mean.transpose(1, 2),
+            "y_log_scale": y_log_scale.transpose(1, 2),
+            "alignments": attn,
+            "durations_log": o_dur_log.transpose(1, 2),
+            "total_durations_log": o_attn_dur.transpose(1, 2),
+        }
+        return outputs
+
+    @torch.no_grad()
+    def inference_with_MAS(
+        self, x, x_lengths, y=None, y_lengths=None, aux_input={"d_vectors": None, "speaker_ids": None}
+    ):  # pylint: disable=dangerous-default-value
+        """
+        It's similar to the teacher forcing in Tacotron.
+        It was proposed in: https://arxiv.org/abs/2104.05557
+
+        Shapes:
+            - x: :math:`[B, T]`
+            - x_lenghts: :math:`B`
+            - y: :math:`[B, T, C]`
+            - y_lengths: :math:`B`
+            - g: :math:`[B, C] or B`
+        """
+        y = y.transpose(1, 2)
+        y_max_length = y.size(2)
+        # norm speaker embeddings
+        g = self._speaker_embedding(aux_input)
+        # embedding pass
+        o_mean, o_log_scale, o_dur_log, x_mask = self.encoder(x, x_lengths, g=g)
+        # drop redisual frames wrt num_squeeze and set y_lengths.
+        y, y_lengths, y_max_length, attn = self.preprocess(y, y_lengths, y_max_length, None)
+        # create masks
+        y_mask = torch.unsqueeze(sequence_mask(y_lengths, y_max_length), 1).to(x_mask.dtype)
+        attn_mask = torch.unsqueeze(x_mask, -1) * torch.unsqueeze(y_mask, 2)
+        # decoder pass
+        z, logdet = self.decoder(y, y_mask, g=g, reverse=False)
+        # find the alignment path between z and encoder output
+        o_scale = torch.exp(-2 * o_log_scale)
+        logp1 = torch.sum(-0.5 * math.log(2 * math.pi) - o_log_scale, [1]).unsqueeze(-1)  # [b, t, 1]
+        logp2 = torch.matmul(o_scale.transpose(1, 2), -0.5 * (z**2))  # [b, t, d] x [b, d, t'] = [b, t, t']
+        logp3 = torch.matmul((o_mean * o_scale).transpose(1, 2), z)  # [b, t, d] x [b, d, t'] = [b, t, t']
+        logp4 = torch.sum(-0.5 * (o_mean**2) * o_scale, [1]).unsqueeze(-1)  # [b, t, 1]
+        logp = logp1 + logp2 + logp3 + logp4  # [b, t, t']
+        attn = maximum_path(logp, attn_mask.squeeze(1)).unsqueeze(1).detach()
+
+        y_mean, y_log_scale, o_attn_dur = self.compute_outputs(attn, o_mean, o_log_scale, x_mask)
+        attn = attn.squeeze(1).permute(0, 2, 1)
+
+        # get predited aligned distribution
+        z = y_mean * y_mask
+
+        # reverse the decoder and predict using the aligned distribution
+        y, logdet = self.decoder(z, y_mask, g=g, reverse=True)
+        outputs = {
+            "model_outputs": z.transpose(1, 2),
+            "logdet": logdet,
+            "y_mean": y_mean.transpose(1, 2),
+            "y_log_scale": y_log_scale.transpose(1, 2),
+            "alignments": attn,
+            "durations_log": o_dur_log.transpose(1, 2),
+            "total_durations_log": o_attn_dur.transpose(1, 2),
+        }
+        return outputs
+
+    @torch.no_grad()
+    def decoder_inference(
+        self, y, y_lengths=None, aux_input={"d_vectors": None, "speaker_ids": None}
+    ):  # pylint: disable=dangerous-default-value
+        """
+        Shapes:
+            - y: :math:`[B, T, C]`
+            - y_lengths: :math:`B`
+            - g: :math:`[B, C] or B`
+        """
+        y = y.transpose(1, 2)
+        y_max_length = y.size(2)
+        g = self._speaker_embedding(aux_input)
+        y_mask = torch.unsqueeze(sequence_mask(y_lengths, y_max_length), 1).to(y.dtype)
+        # decoder pass
+        z, logdet = self.decoder(y, y_mask, g=g, reverse=False)
+        # reverse decoder and predict
+        y, logdet = self.decoder(z, y_mask, g=g, reverse=True)
+        outputs = {}
+        outputs["model_outputs"] = y.transpose(1, 2)
+        outputs["logdet"] = logdet
+        return outputs
+
+    @torch.no_grad()
+    def inference(
+        self, x, aux_input={"x_lengths": None, "d_vectors": None, "speaker_ids": None}
+    ):  # pylint: disable=dangerous-default-value
+        x_lengths = aux_input["x_lengths"]
+        g = self._speaker_embedding(aux_input)
+        # embedding pass
+        o_mean, o_log_scale, o_dur_log, x_mask = self.encoder(x, x_lengths, g=g)
+        # compute output durations
+        w = (torch.exp(o_dur_log) - 1) * x_mask * self.length_scale
+        w_ceil = torch.clamp_min(torch.ceil(w), 1)
+        y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
+        y_max_length = None
+        # compute masks
+        y_mask = torch.unsqueeze(sequence_mask(y_lengths, y_max_length), 1).to(x_mask.dtype)
+        attn_mask = torch.unsqueeze(x_mask, -1) * torch.unsqueeze(y_mask, 2)
+        # compute attention mask
+        attn = generate_path(w_ceil.squeeze(1), attn_mask.squeeze(1)).unsqueeze(1)
+        y_mean, y_log_scale, o_attn_dur = self.compute_outputs(attn, o_mean, o_log_scale, x_mask)
+
+        z = (y_mean + torch.exp(y_log_scale) * torch.randn_like(y_mean) * self.inference_noise_scale) * y_mask
+        # decoder pass
+        y, logdet = self.decoder(z, y_mask, g=g, reverse=True)
+        attn = attn.squeeze(1).permute(0, 2, 1)
+        outputs = {
+            "model_outputs": y.transpose(1, 2),
+            "logdet": logdet,
+            "y_mean": y_mean.transpose(1, 2),
+            "y_log_scale": y_log_scale.transpose(1, 2),
+            "alignments": attn,
+            "durations_log": o_dur_log.transpose(1, 2),
+            "total_durations_log": o_attn_dur.transpose(1, 2),
+        }
+        return outputs
+
+    def train_step(self, batch: dict, criterion: nn.Module):
+        """A single training step. Forward pass and loss computation. Run data depended initialization for the
+        first `config.data_dep_init_steps` steps.
+
+        Args:
+            batch (dict): [description]
+            criterion (nn.Module): [description]
+        """
+        text_input = batch["text_input"]
+        text_lengths = batch["text_lengths"]
+        mel_input = batch["mel_input"]
+        mel_lengths = batch["mel_lengths"]
+        d_vectors = batch["d_vectors"]
+        speaker_ids = batch["speaker_ids"]
+
+        if self.run_data_dep_init and self.training:
+            # compute data-dependent initialization of activation norm layers
+            self.unlock_act_norm_layers()
+            with torch.no_grad():
+                _ = self.forward(
+                    text_input,
+                    text_lengths,
+                    mel_input,
+                    mel_lengths,
+                    aux_input={"d_vectors": d_vectors, "speaker_ids": speaker_ids},
+                )
+            outputs = None
+            loss_dict = None
+            self.lock_act_norm_layers()
+        else:
+            # normal training step
+            outputs = self.forward(
+                text_input,
+                text_lengths,
+                mel_input,
+                mel_lengths,
+                aux_input={"d_vectors": d_vectors, "speaker_ids": speaker_ids},
+            )
+
+            with autocast(enabled=False):  # avoid mixed_precision in criterion
+                loss_dict = criterion(
+                    outputs["z"].float(),
+                    outputs["y_mean"].float(),
+                    outputs["y_log_scale"].float(),
+                    outputs["logdet"].float(),
+                    mel_lengths,
+                    outputs["durations_log"].float(),
+                    outputs["total_durations_log"].float(),
+                    text_lengths,
+                )
+        return outputs, loss_dict
+
+    def _create_logs(self, batch, outputs, ap):
+        alignments = outputs["alignments"]
+        text_input = batch["text_input"][:1] if batch["text_input"] is not None else None
+        text_lengths = batch["text_lengths"]
+        mel_input = batch["mel_input"]
+        d_vectors = batch["d_vectors"][:1] if batch["d_vectors"] is not None else None
+        speaker_ids = batch["speaker_ids"][:1] if batch["speaker_ids"] is not None else None
+
+        # model runs reverse flow to predict spectrograms
+        pred_outputs = self.inference(
+            text_input,
+            aux_input={"x_lengths": text_lengths[:1], "d_vectors": d_vectors, "speaker_ids": speaker_ids},
+        )
+        model_outputs = pred_outputs["model_outputs"]
+
+        pred_spec = model_outputs[0].data.cpu().numpy()
+        gt_spec = mel_input[0].data.cpu().numpy()
+        align_img = alignments[0].data.cpu().numpy()
+
+        figures = {
+            "prediction": plot_spectrogram(pred_spec, ap, output_fig=False),
+            "ground_truth": plot_spectrogram(gt_spec, ap, output_fig=False),
+            "alignment": plot_alignment(align_img, output_fig=False),
+        }
+
+        # Sample audio
+        train_audio = ap.inv_melspectrogram(pred_spec.T)
+        return figures, {"audio": train_audio}
+
+    def train_log(
+        self, batch: dict, outputs: dict, logger: "Logger", assets: dict, steps: int
+    ) -> None:  # pylint: disable=no-self-use
+        figures, audios = self._create_logs(batch, outputs, self.ap)
+        logger.train_figures(steps, figures)
+        logger.train_audios(steps, audios, self.ap.sample_rate)
+
+    @torch.no_grad()
+    def eval_step(self, batch: dict, criterion: nn.Module):
+        return self.train_step(batch, criterion)
+
+    def eval_log(self, batch: dict, outputs: dict, logger: "Logger", assets: dict, steps: int) -> None:
+        figures, audios = self._create_logs(batch, outputs, self.ap)
+        logger.eval_figures(steps, figures)
+        logger.eval_audios(steps, audios, self.ap.sample_rate)
+
+    @torch.no_grad()
+    def test_run(self, assets: Dict) -> Tuple[Dict, Dict]:
+        """Generic test run for `tts` models used by `Trainer`.
+
+        You can override this for a different behaviour.
+
+        Returns:
+            Tuple[Dict, Dict]: Test figures and audios to be projected to Tensorboard.
+        """
+        print(" | > Synthesizing test sentences.")
+        test_audios = {}
+        test_figures = {}
+        test_sentences = self.config.test_sentences
+        aux_inputs = self._get_test_aux_input()
+        if len(test_sentences) == 0:
+            print(" | [!] No test sentences provided.")
+        else:
+            for idx, sen in enumerate(test_sentences):
+                outputs = synthesis(
+                    self,
+                    sen,
+                    self.config,
+                    "cuda" in str(next(self.parameters()).device),
+                    speaker_id=aux_inputs["speaker_id"],
+                    d_vector=aux_inputs["d_vector"],
+                    style_wav=aux_inputs["style_wav"],
+                    use_griffin_lim=True,
+                    do_trim_silence=False,
+                )
+
+                test_audios["{}-audio".format(idx)] = outputs["wav"]
+                test_figures["{}-prediction".format(idx)] = plot_spectrogram(
+                    outputs["outputs"]["model_outputs"], self.ap, output_fig=False
+                )
+                test_figures["{}-alignment".format(idx)] = plot_alignment(outputs["alignments"], output_fig=False)
+        return test_figures, test_audios
+
+    def preprocess(self, y, y_lengths, y_max_length, attn=None):
+        if y_max_length is not None:
+            y_max_length = (y_max_length // self.num_squeeze) * self.num_squeeze
+            y = y[:, :, :y_max_length]
+            if attn is not None:
+                attn = attn[:, :, :, :y_max_length]
+        y_lengths = torch.div(y_lengths, self.num_squeeze, rounding_mode="floor") * self.num_squeeze
+        return y, y_lengths, y_max_length, attn
+
+    def store_inverse(self):
+        self.decoder.store_inverse()
+
+    def load_checkpoint(
+        self, config, checkpoint_path, eval=False
+    ):  # pylint: disable=unused-argument, redefined-builtin
+        state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"))
+        self.load_state_dict(state["model"])
+        if eval:
+            self.eval()
+            self.store_inverse()
+            assert not self.training
+
+    @staticmethod
+    def get_criterion():
+        from TTS.tts.layers.losses import GlowTTSLoss  # pylint: disable=import-outside-toplevel
+
+        return GlowTTSLoss()
+
+    def on_train_step_start(self, trainer):
+        """Decide on every training step wheter enable/disable data depended initialization."""
+        self.run_data_dep_init = trainer.total_steps_done < self.data_dep_init_steps
+
+    @staticmethod
+    def init_from_config(config: "GlowTTSConfig", samples: Union[List[List], List[Dict]] = None, verbose=True):
+        """Initiate model from config
+
+        Args:
+            config (VitsConfig): Model config.
+            samples (Union[List[List], List[Dict]]): Training samples to parse speaker ids for training.
+                Defaults to None.
+            verbose (bool): If True, print init messages. Defaults to True.
+        """
+        from TTS.utils.audio import AudioProcessor
+
+        ap = AudioProcessor.init_from_config(config, verbose)
+        tokenizer, new_config = TTSTokenizer.init_from_config(config)
+        speaker_manager = SpeakerManager.init_from_config(config, samples)
+        return GlowTTS(new_config, ap, tokenizer, speaker_manager)

TTS/tts/models/neuralhmm_tts.py

@@ -0,0 +1,385 @@
+import os
+from typing import Dict, List, Union
+
+import torch
+from coqpit import Coqpit
+from torch import nn
+from trainer.logging.tensorboard_logger import TensorboardLogger
+
+from TTS.tts.layers.overflow.common_layers import Encoder, OverflowUtils
+from TTS.tts.layers.overflow.neural_hmm import NeuralHMM
+from TTS.tts.layers.overflow.plotting_utils import (
+    get_spec_from_most_probable_state,
+    plot_transition_probabilities_to_numpy,
+)
+from TTS.tts.models.base_tts import BaseTTS
+from TTS.tts.utils.speakers import SpeakerManager
+from TTS.tts.utils.text.tokenizer import TTSTokenizer
+from TTS.tts.utils.visual import plot_alignment, plot_spectrogram
+from TTS.utils.generic_utils import format_aux_input
+from TTS.utils.io import load_fsspec
+
+
+class NeuralhmmTTS(BaseTTS):
+    """Neural HMM TTS model.
+
+    Paper::
+        https://arxiv.org/abs/2108.13320
+
+    Paper abstract::
+        Neural sequence-to-sequence TTS has achieved significantly better output quality
+    than statistical speech synthesis using HMMs.However, neural TTS is generally not probabilistic
+    and uses non-monotonic attention. Attention failures increase training time and can make
+    synthesis babble incoherently. This paper describes how the old and new paradigms can be
+    combined to obtain the advantages of both worlds, by replacing attention in neural TTS with
+    an autoregressive left-right no-skip hidden Markov model defined by a neural network.
+    Based on this proposal, we modify Tacotron 2 to obtain an HMM-based neural TTS model with
+    monotonic alignment, trained to maximise the full sequence likelihood without approximation.
+    We also describe how to combine ideas from classical and contemporary TTS for best results.
+    The resulting example system is smaller and simpler than Tacotron 2, and learns to speak with
+    fewer iterations and less data, whilst achieving comparable naturalness prior to the post-net.
+    Our approach also allows easy control over speaking rate. Audio examples and code
+    are available at https://shivammehta25.github.io/Neural-HMM/ .
+
+    Note:
+        - This is a parameter efficient version of OverFlow (15.3M vs 28.6M). Since it has half the
+        number of parameters as OverFlow the synthesis output quality is suboptimal (but comparable to Tacotron2
+        without Postnet), but it learns to speak with even lesser amount of data and is still significantly faster
+        than other attention-based methods.
+
+        - Neural HMMs uses flat start initialization i.e it computes the means and std and transition probabilities
+        of the dataset and uses them to initialize the model. This benefits the model and helps with faster learning
+        If you change the dataset or want to regenerate the parameters change the `force_generate_statistics` and
+        `mel_statistics_parameter_path` accordingly.
+
+        - To enable multi-GPU training, set the `use_grad_checkpointing=False` in config.
+        This will significantly increase the memory usage.  This is because to compute
+        the actual data likelihood (not an approximation using MAS/Viterbi) we must use
+        all the states at the previous time step during the forward pass to decide the
+        probability distribution at the current step i.e the difference between the forward
+        algorithm and viterbi approximation.
+
+    Check :class:`TTS.tts.configs.neuralhmm_tts_config.NeuralhmmTTSConfig` for class arguments.
+    """
+
+    def __init__(
+        self,
+        config: "NeuralhmmTTSConfig",
+        ap: "AudioProcessor" = None,
+        tokenizer: "TTSTokenizer" = None,
+        speaker_manager: SpeakerManager = None,
+    ):
+        super().__init__(config, ap, tokenizer, speaker_manager)
+
+        # pass all config fields to `self`
+        # for fewer code change
+        self.config = config
+        for key in config:
+            setattr(self, key, config[key])
+
+        self.encoder = Encoder(config.num_chars, config.state_per_phone, config.encoder_in_out_features)
+        self.neural_hmm = NeuralHMM(
+            frame_channels=self.out_channels,
+            ar_order=self.ar_order,
+            deterministic_transition=self.deterministic_transition,
+            encoder_dim=self.encoder_in_out_features,
+            prenet_type=self.prenet_type,
+            prenet_dim=self.prenet_dim,
+            prenet_n_layers=self.prenet_n_layers,
+            prenet_dropout=self.prenet_dropout,
+            prenet_dropout_at_inference=self.prenet_dropout_at_inference,
+            memory_rnn_dim=self.memory_rnn_dim,
+            outputnet_size=self.outputnet_size,
+            flat_start_params=self.flat_start_params,
+            std_floor=self.std_floor,
+            use_grad_checkpointing=self.use_grad_checkpointing,
+        )
+
+        self.register_buffer("mean", torch.tensor(0))
+        self.register_buffer("std", torch.tensor(1))
+
+    def update_mean_std(self, statistics_dict: Dict):
+        self.mean.data = torch.tensor(statistics_dict["mean"])
+        self.std.data = torch.tensor(statistics_dict["std"])
+
+    def preprocess_batch(self, text, text_len, mels, mel_len):
+        if self.mean.item() == 0 or self.std.item() == 1:
+            statistics_dict = torch.load(self.mel_statistics_parameter_path)
+            self.update_mean_std(statistics_dict)
+
+        mels = self.normalize(mels)
+        return text, text_len, mels, mel_len
+
+    def normalize(self, x):
+        return x.sub(self.mean).div(self.std)
+
+    def inverse_normalize(self, x):
+        return x.mul(self.std).add(self.mean)
+
+    def forward(self, text, text_len, mels, mel_len):
+        """
+        Forward pass for training and computing the log likelihood of a given batch.
+
+        Shapes:
+            Shapes:
+            text: :math:`[B, T_in]`
+            text_len: :math:`[B]`
+            mels: :math:`[B, T_out, C]`
+            mel_len: :math:`[B]`
+        """
+        text, text_len, mels, mel_len = self.preprocess_batch(text, text_len, mels, mel_len)
+        encoder_outputs, encoder_output_len = self.encoder(text, text_len)
+
+        log_probs, fwd_alignments, transition_vectors, means = self.neural_hmm(
+            encoder_outputs, encoder_output_len, mels.transpose(1, 2), mel_len
+        )
+
+        outputs = {
+            "log_probs": log_probs,
+            "alignments": fwd_alignments,
+            "transition_vectors": transition_vectors,
+            "means": means,
+        }
+
+        return outputs
+
+    @staticmethod
+    def _training_stats(batch):
+        stats = {}
+        stats["avg_text_length"] = batch["text_lengths"].float().mean()
+        stats["avg_spec_length"] = batch["mel_lengths"].float().mean()
+        stats["avg_text_batch_occupancy"] = (batch["text_lengths"].float() / batch["text_lengths"].float().max()).mean()
+        stats["avg_spec_batch_occupancy"] = (batch["mel_lengths"].float() / batch["mel_lengths"].float().max()).mean()
+        return stats
+
+    def train_step(self, batch: dict, criterion: nn.Module):
+        text_input = batch["text_input"]
+        text_lengths = batch["text_lengths"]
+        mel_input = batch["mel_input"]
+        mel_lengths = batch["mel_lengths"]
+
+        outputs = self.forward(
+            text=text_input,
+            text_len=text_lengths,
+            mels=mel_input,
+            mel_len=mel_lengths,
+        )
+        loss_dict = criterion(outputs["log_probs"] / (mel_lengths.sum() + text_lengths.sum()))
+
+        # for printing useful statistics on terminal
+        loss_dict.update(self._training_stats(batch))
+        return outputs, loss_dict
+
+    def eval_step(self, batch: Dict, criterion: nn.Module):
+        return self.train_step(batch, criterion)
+
+    def _format_aux_input(self, aux_input: Dict, default_input_dict):
+        """Set missing fields to their default value.
+
+        Args:
+            aux_inputs (Dict): Dictionary containing the auxiliary inputs.
+        """
+        default_input_dict = default_input_dict.copy()
+        default_input_dict.update(
+            {
+                "sampling_temp": self.sampling_temp,
+                "max_sampling_time": self.max_sampling_time,
+                "duration_threshold": self.duration_threshold,
+            }
+        )
+        if aux_input:
+            return format_aux_input(default_input_dict, aux_input)
+        return default_input_dict
+
+    @torch.no_grad()
+    def inference(
+        self,
+        text: torch.Tensor,
+        aux_input={"x_lengths": None, "sampling_temp": None, "max_sampling_time": None, "duration_threshold": None},
+    ):  # pylint: disable=dangerous-default-value
+        """Sampling from the model
+
+        Args:
+            text (torch.Tensor): :math:`[B, T_in]`
+            aux_inputs (_type_, optional): _description_. Defaults to None.
+
+        Returns:
+            outputs: Dictionary containing the following
+                - mel (torch.Tensor): :math:`[B, T_out, C]`
+                - hmm_outputs_len (torch.Tensor): :math:`[B]`
+                - state_travelled (List[List[int]]): List of lists containing the state travelled for each sample in the batch.
+                - input_parameters (list[torch.FloatTensor]): Input parameters to the neural HMM.
+                - output_parameters (list[torch.FloatTensor]): Output parameters to the neural HMM.
+        """
+        default_input_dict = {
+            "x_lengths": torch.sum(text != 0, dim=1),
+        }
+        aux_input = self._format_aux_input(aux_input, default_input_dict)
+        encoder_outputs, encoder_output_len = self.encoder.inference(text, aux_input["x_lengths"])
+        outputs = self.neural_hmm.inference(
+            encoder_outputs,
+            encoder_output_len,
+            sampling_temp=aux_input["sampling_temp"],
+            max_sampling_time=aux_input["max_sampling_time"],
+            duration_threshold=aux_input["duration_threshold"],
+        )
+        mels, mel_outputs_len = outputs["hmm_outputs"], outputs["hmm_outputs_len"]
+
+        mels = self.inverse_normalize(mels)
+        outputs.update({"model_outputs": mels, "model_outputs_len": mel_outputs_len})
+        outputs["alignments"] = OverflowUtils.double_pad(outputs["alignments"])
+        return outputs
+
+    @staticmethod
+    def get_criterion():
+        return NLLLoss()
+
+    @staticmethod
+    def init_from_config(config: "NeuralhmmTTSConfig", samples: Union[List[List], List[Dict]] = None, verbose=True):
+        """Initiate model from config
+
+        Args:
+            config (VitsConfig): Model config.
+            samples (Union[List[List], List[Dict]]): Training samples to parse speaker ids for training.
+                Defaults to None.
+            verbose (bool): If True, print init messages. Defaults to True.
+        """
+        from TTS.utils.audio import AudioProcessor
+
+        ap = AudioProcessor.init_from_config(config, verbose)
+        tokenizer, new_config = TTSTokenizer.init_from_config(config)
+        speaker_manager = SpeakerManager.init_from_config(config, samples)
+        return NeuralhmmTTS(new_config, ap, tokenizer, speaker_manager)
+
+    def load_checkpoint(
+        self, config: Coqpit, checkpoint_path: str, eval: bool = False, strict: bool = True, cache=False
+    ):  # pylint: disable=unused-argument, redefined-builtin
+        state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"))
+        self.load_state_dict(state["model"])
+        if eval:
+            self.eval()
+            assert not self.training
+
+    def on_init_start(self, trainer):
+        """If the current dataset does not have normalisation statistics and initialisation transition_probability it computes them otherwise loads."""
+        if not os.path.isfile(trainer.config.mel_statistics_parameter_path) or trainer.config.force_generate_statistics:
+            dataloader = trainer.get_train_dataloader(
+                training_assets=None, samples=trainer.train_samples, verbose=False
+            )
+            print(
+                f" | > Data parameters not found for: {trainer.config.mel_statistics_parameter_path}. Computing mel normalization parameters..."
+            )
+            data_mean, data_std, init_transition_prob = OverflowUtils.get_data_parameters_for_flat_start(
+                dataloader, trainer.config.out_channels, trainer.config.state_per_phone
+            )
+            print(
+                f" | > Saving data parameters to: {trainer.config.mel_statistics_parameter_path}: value: {data_mean, data_std, init_transition_prob}"
+            )
+            statistics = {
+                "mean": data_mean.item(),
+                "std": data_std.item(),
+                "init_transition_prob": init_transition_prob.item(),
+            }
+            torch.save(statistics, trainer.config.mel_statistics_parameter_path)
+
+        else:
+            print(
+                f" | > Data parameters found for: {trainer.config.mel_statistics_parameter_path}. Loading mel normalization parameters..."
+            )
+            statistics = torch.load(trainer.config.mel_statistics_parameter_path)
+            data_mean, data_std, init_transition_prob = (
+                statistics["mean"],
+                statistics["std"],
+                statistics["init_transition_prob"],
+            )
+            print(f" | > Data parameters loaded with value: {data_mean, data_std, init_transition_prob}")
+
+        trainer.config.flat_start_params["transition_p"] = (
+            init_transition_prob.item() if torch.is_tensor(init_transition_prob) else init_transition_prob
+        )
+        OverflowUtils.update_flat_start_transition(trainer.model, init_transition_prob)
+        trainer.model.update_mean_std(statistics)
+
+    @torch.inference_mode()
+    def _create_logs(self, batch, outputs, ap):  # pylint: disable=no-self-use, unused-argument
+        alignments, transition_vectors = outputs["alignments"], outputs["transition_vectors"]
+        means = torch.stack(outputs["means"], dim=1)
+
+        figures = {
+            "alignment": plot_alignment(alignments[0].exp(), title="Forward alignment", fig_size=(20, 20)),
+            "log_alignment": plot_alignment(
+                alignments[0].exp(), title="Forward log alignment", plot_log=True, fig_size=(20, 20)
+            ),
+            "transition_vectors": plot_alignment(transition_vectors[0], title="Transition vectors", fig_size=(20, 20)),
+            "mel_from_most_probable_state": plot_spectrogram(
+                get_spec_from_most_probable_state(alignments[0], means[0]), fig_size=(12, 3)
+            ),
+            "mel_target": plot_spectrogram(batch["mel_input"][0], fig_size=(12, 3)),
+        }
+
+        # sample one item from the batch -1 will give the smalles item
+        print(" | > Synthesising audio from the model...")
+        inference_output = self.inference(
+            batch["text_input"][-1].unsqueeze(0), aux_input={"x_lengths": batch["text_lengths"][-1].unsqueeze(0)}
+        )
+        figures["synthesised"] = plot_spectrogram(inference_output["model_outputs"][0], fig_size=(12, 3))
+
+        states = [p[1] for p in inference_output["input_parameters"][0]]
+        transition_probability_synthesising = [p[2].cpu().numpy() for p in inference_output["output_parameters"][0]]
+
+        for i in range((len(transition_probability_synthesising) // 200) + 1):
+            start = i * 200
+            end = (i + 1) * 200
+            figures[f"synthesised_transition_probabilities/{i}"] = plot_transition_probabilities_to_numpy(
+                states[start:end], transition_probability_synthesising[start:end]
+            )
+
+        audio = ap.inv_melspectrogram(inference_output["model_outputs"][0].T.cpu().numpy())
+        return figures, {"audios": audio}
+
+    def train_log(
+        self, batch: dict, outputs: dict, logger: "Logger", assets: dict, steps: int
+    ):  # pylint: disable=unused-argument
+        """Log training progress."""
+        figures, audios = self._create_logs(batch, outputs, self.ap)
+        logger.train_figures(steps, figures)
+        logger.train_audios(steps, audios, self.ap.sample_rate)
+
+    def eval_log(
+        self, batch: Dict, outputs: Dict, logger: "Logger", assets: Dict, steps: int
+    ):  # pylint: disable=unused-argument
+        """Compute and log evaluation metrics."""
+        # Plot model parameters histograms
+        if isinstance(logger, TensorboardLogger):
+            # I don't know if any other loggers supports this
+            for tag, value in self.named_parameters():
+                tag = tag.replace(".", "/")
+                logger.writer.add_histogram(tag, value.data.cpu().numpy(), steps)
+
+        figures, audios = self._create_logs(batch, outputs, self.ap)
+        logger.eval_figures(steps, figures)
+        logger.eval_audios(steps, audios, self.ap.sample_rate)
+
+    def test_log(
+        self, outputs: dict, logger: "Logger", assets: dict, steps: int  # pylint: disable=unused-argument
+    ) -> None:
+        logger.test_audios(steps, outputs[1], self.ap.sample_rate)
+        logger.test_figures(steps, outputs[0])
+
+
+class NLLLoss(nn.Module):
+    """Negative log likelihood loss."""
+
+    def forward(self, log_prob: torch.Tensor) -> dict:  # pylint: disable=no-self-use
+        """Compute the loss.
+
+        Args:
+            logits (Tensor): [B, T, D]
+
+        Returns:
+            Tensor: [1]
+
+        """
+        return_dict = {}
+        return_dict["loss"] = -log_prob.mean()
+        return return_dict

TTS/tts/models/overflow.py

@@ -0,0 +1,401 @@
+import os
+from typing import Dict, List, Union
+
+import torch
+from coqpit import Coqpit
+from torch import nn
+from trainer.logging.tensorboard_logger import TensorboardLogger
+
+from TTS.tts.layers.overflow.common_layers import Encoder, OverflowUtils
+from TTS.tts.layers.overflow.decoder import Decoder
+from TTS.tts.layers.overflow.neural_hmm import NeuralHMM
+from TTS.tts.layers.overflow.plotting_utils import (
+    get_spec_from_most_probable_state,
+    plot_transition_probabilities_to_numpy,
+)
+from TTS.tts.models.base_tts import BaseTTS
+from TTS.tts.utils.speakers import SpeakerManager
+from TTS.tts.utils.text.tokenizer import TTSTokenizer
+from TTS.tts.utils.visual import plot_alignment, plot_spectrogram
+from TTS.utils.generic_utils import format_aux_input
+from TTS.utils.io import load_fsspec
+
+
+class Overflow(BaseTTS):
+    """OverFlow TTS model.
+
+    Paper::
+        https://arxiv.org/abs/2211.06892
+
+    Paper abstract::
+        Neural HMMs are a type of neural transducer recently proposed for
+    sequence-to-sequence modelling in text-to-speech. They combine the best features
+    of classic statistical speech synthesis and modern neural TTS, requiring less
+    data and fewer training updates, and are less prone to gibberish output caused
+    by neural attention failures. In this paper, we combine neural HMM TTS with
+    normalising flows for describing the highly non-Gaussian distribution of speech
+    acoustics. The result is a powerful, fully probabilistic model of durations and
+    acoustics that can be trained using exact maximum likelihood. Compared to
+    dominant flow-based acoustic models, our approach integrates autoregression for
+    improved modelling of long-range dependences such as utterance-level prosody.
+    Experiments show that a system based on our proposal gives more accurate
+    pronunciations and better subjective speech quality than comparable methods,
+    whilst retaining the original advantages of neural HMMs. Audio examples and code
+    are available at https://shivammehta25.github.io/OverFlow/.
+
+    Note:
+        - Neural HMMs uses flat start initialization i.e it computes the means and std and transition probabilities
+        of the dataset and uses them to initialize the model. This benefits the model and helps with faster learning
+        If you change the dataset or want to regenerate the parameters change the `force_generate_statistics` and
+        `mel_statistics_parameter_path` accordingly.
+
+        - To enable multi-GPU training, set the `use_grad_checkpointing=False` in config.
+        This will significantly increase the memory usage.  This is because to compute
+        the actual data likelihood (not an approximation using MAS/Viterbi) we must use
+        all the states at the previous time step during the forward pass to decide the
+        probability distribution at the current step i.e the difference between the forward
+        algorithm and viterbi approximation.
+
+    Check :class:`TTS.tts.configs.overflow.OverFlowConfig` for class arguments.
+    """
+
+    def __init__(
+        self,
+        config: "OverFlowConfig",
+        ap: "AudioProcessor" = None,
+        tokenizer: "TTSTokenizer" = None,
+        speaker_manager: SpeakerManager = None,
+    ):
+        super().__init__(config, ap, tokenizer, speaker_manager)
+
+        # pass all config fields to `self`
+        # for fewer code change
+        self.config = config
+        for key in config:
+            setattr(self, key, config[key])
+
+        self.decoder_output_dim = config.out_channels
+
+        self.encoder = Encoder(config.num_chars, config.state_per_phone, config.encoder_in_out_features)
+        self.neural_hmm = NeuralHMM(
+            frame_channels=self.out_channels,
+            ar_order=self.ar_order,
+            deterministic_transition=self.deterministic_transition,
+            encoder_dim=self.encoder_in_out_features,
+            prenet_type=self.prenet_type,
+            prenet_dim=self.prenet_dim,
+            prenet_n_layers=self.prenet_n_layers,
+            prenet_dropout=self.prenet_dropout,
+            prenet_dropout_at_inference=self.prenet_dropout_at_inference,
+            memory_rnn_dim=self.memory_rnn_dim,
+            outputnet_size=self.outputnet_size,
+            flat_start_params=self.flat_start_params,
+            std_floor=self.std_floor,
+            use_grad_checkpointing=self.use_grad_checkpointing,
+        )
+
+        self.decoder = Decoder(
+            self.out_channels,
+            self.hidden_channels_dec,
+            self.kernel_size_dec,
+            self.dilation_rate,
+            self.num_flow_blocks_dec,
+            self.num_block_layers,
+            dropout_p=self.dropout_p_dec,
+            num_splits=self.num_splits,
+            num_squeeze=self.num_squeeze,
+            sigmoid_scale=self.sigmoid_scale,
+            c_in_channels=self.c_in_channels,
+        )
+
+        self.register_buffer("mean", torch.tensor(0))
+        self.register_buffer("std", torch.tensor(1))
+
+    def update_mean_std(self, statistics_dict: Dict):
+        self.mean.data = torch.tensor(statistics_dict["mean"])
+        self.std.data = torch.tensor(statistics_dict["std"])
+
+    def preprocess_batch(self, text, text_len, mels, mel_len):
+        if self.mean.item() == 0 or self.std.item() == 1:
+            statistics_dict = torch.load(self.mel_statistics_parameter_path)
+            self.update_mean_std(statistics_dict)
+
+        mels = self.normalize(mels)
+        return text, text_len, mels, mel_len
+
+    def normalize(self, x):
+        return x.sub(self.mean).div(self.std)
+
+    def inverse_normalize(self, x):
+        return x.mul(self.std).add(self.mean)
+
+    def forward(self, text, text_len, mels, mel_len):
+        """
+        Forward pass for training and computing the log likelihood of a given batch.
+
+        Shapes:
+            Shapes:
+            text: :math:`[B, T_in]`
+            text_len: :math:`[B]`
+            mels: :math:`[B, T_out, C]`
+            mel_len: :math:`[B]`
+        """
+        text, text_len, mels, mel_len = self.preprocess_batch(text, text_len, mels, mel_len)
+        encoder_outputs, encoder_output_len = self.encoder(text, text_len)
+        z, z_lengths, logdet = self.decoder(mels.transpose(1, 2), mel_len)
+        log_probs, fwd_alignments, transition_vectors, means = self.neural_hmm(
+            encoder_outputs, encoder_output_len, z, z_lengths
+        )
+
+        outputs = {
+            "log_probs": log_probs + logdet,
+            "alignments": fwd_alignments,
+            "transition_vectors": transition_vectors,
+            "means": means,
+        }
+
+        return outputs
+
+    @staticmethod
+    def _training_stats(batch):
+        stats = {}
+        stats["avg_text_length"] = batch["text_lengths"].float().mean()
+        stats["avg_spec_length"] = batch["mel_lengths"].float().mean()
+        stats["avg_text_batch_occupancy"] = (batch["text_lengths"].float() / batch["text_lengths"].float().max()).mean()
+        stats["avg_spec_batch_occupancy"] = (batch["mel_lengths"].float() / batch["mel_lengths"].float().max()).mean()
+        return stats
+
+    def train_step(self, batch: dict, criterion: nn.Module):
+        text_input = batch["text_input"]
+        text_lengths = batch["text_lengths"]
+        mel_input = batch["mel_input"]
+        mel_lengths = batch["mel_lengths"]
+
+        outputs = self.forward(
+            text=text_input,
+            text_len=text_lengths,
+            mels=mel_input,
+            mel_len=mel_lengths,
+        )
+        loss_dict = criterion(outputs["log_probs"] / (mel_lengths.sum() + text_lengths.sum()))
+
+        # for printing useful statistics on terminal
+        loss_dict.update(self._training_stats(batch))
+        return outputs, loss_dict
+
+    def eval_step(self, batch: Dict, criterion: nn.Module):
+        return self.train_step(batch, criterion)
+
+    def _format_aux_input(self, aux_input: Dict, default_input_dict):
+        """Set missing fields to their default value.
+
+        Args:
+            aux_inputs (Dict): Dictionary containing the auxiliary inputs.
+        """
+        default_input_dict = default_input_dict.copy()
+        default_input_dict.update(
+            {
+                "sampling_temp": self.sampling_temp,
+                "max_sampling_time": self.max_sampling_time,
+                "duration_threshold": self.duration_threshold,
+            }
+        )
+        if aux_input:
+            return format_aux_input(default_input_dict, aux_input)
+        return default_input_dict
+
+    @torch.no_grad()
+    def inference(
+        self,
+        text: torch.Tensor,
+        aux_input={"x_lengths": None, "sampling_temp": None, "max_sampling_time": None, "duration_threshold": None},
+    ):  # pylint: disable=dangerous-default-value
+        """Sampling from the model
+
+        Args:
+            text (torch.Tensor): :math:`[B, T_in]`
+            aux_inputs (_type_, optional): _description_. Defaults to None.
+
+        Returns:
+            outputs: Dictionary containing the following
+                - mel (torch.Tensor): :math:`[B, T_out, C]`
+                - hmm_outputs_len (torch.Tensor): :math:`[B]`
+                - state_travelled (List[List[int]]): List of lists containing the state travelled for each sample in the batch.
+                - input_parameters (list[torch.FloatTensor]): Input parameters to the neural HMM.
+                - output_parameters (list[torch.FloatTensor]): Output parameters to the neural HMM.
+        """
+        default_input_dict = {
+            "x_lengths": torch.sum(text != 0, dim=1),
+        }
+        aux_input = self._format_aux_input(aux_input, default_input_dict)
+        encoder_outputs, encoder_output_len = self.encoder.inference(text, aux_input["x_lengths"])
+        outputs = self.neural_hmm.inference(
+            encoder_outputs,
+            encoder_output_len,
+            sampling_temp=aux_input["sampling_temp"],
+            max_sampling_time=aux_input["max_sampling_time"],
+            duration_threshold=aux_input["duration_threshold"],
+        )
+
+        mels, mel_outputs_len, _ = self.decoder(
+            outputs["hmm_outputs"].transpose(1, 2), outputs["hmm_outputs_len"], reverse=True
+        )
+        mels = self.inverse_normalize(mels.transpose(1, 2))
+        outputs.update({"model_outputs": mels, "model_outputs_len": mel_outputs_len})
+        outputs["alignments"] = OverflowUtils.double_pad(outputs["alignments"])
+        return outputs
+
+    @staticmethod
+    def get_criterion():
+        return NLLLoss()
+
+    @staticmethod
+    def init_from_config(config: "OverFlowConfig", samples: Union[List[List], List[Dict]] = None, verbose=True):
+        """Initiate model from config
+
+        Args:
+            config (VitsConfig): Model config.
+            samples (Union[List[List], List[Dict]]): Training samples to parse speaker ids for training.
+                Defaults to None.
+            verbose (bool): If True, print init messages. Defaults to True.
+        """
+        from TTS.utils.audio import AudioProcessor
+
+        ap = AudioProcessor.init_from_config(config, verbose)
+        tokenizer, new_config = TTSTokenizer.init_from_config(config)
+        speaker_manager = SpeakerManager.init_from_config(config, samples)
+        return Overflow(new_config, ap, tokenizer, speaker_manager)
+
+    def load_checkpoint(
+        self, config: Coqpit, checkpoint_path: str, eval: bool = False, strict: bool = True, cache=False
+    ):  # pylint: disable=unused-argument, redefined-builtin
+        state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"))
+        self.load_state_dict(state["model"])
+        if eval:
+            self.eval()
+            self.decoder.store_inverse()
+            assert not self.training
+
+    def on_init_start(self, trainer):
+        """If the current dataset does not have normalisation statistics and initialisation transition_probability it computes them otherwise loads."""
+        if not os.path.isfile(trainer.config.mel_statistics_parameter_path) or trainer.config.force_generate_statistics:
+            dataloader = trainer.get_train_dataloader(
+                training_assets=None, samples=trainer.train_samples, verbose=False
+            )
+            print(
+                f" | > Data parameters not found for: {trainer.config.mel_statistics_parameter_path}. Computing mel normalization parameters..."
+            )
+            data_mean, data_std, init_transition_prob = OverflowUtils.get_data_parameters_for_flat_start(
+                dataloader, trainer.config.out_channels, trainer.config.state_per_phone
+            )
+            print(
+                f" | > Saving data parameters to: {trainer.config.mel_statistics_parameter_path}: value: {data_mean, data_std, init_transition_prob}"
+            )
+            statistics = {
+                "mean": data_mean.item(),
+                "std": data_std.item(),
+                "init_transition_prob": init_transition_prob.item(),
+            }
+            torch.save(statistics, trainer.config.mel_statistics_parameter_path)
+
+        else:
+            print(
+                f" | > Data parameters found for: {trainer.config.mel_statistics_parameter_path}. Loading mel normalization parameters..."
+            )
+            statistics = torch.load(trainer.config.mel_statistics_parameter_path)
+            data_mean, data_std, init_transition_prob = (
+                statistics["mean"],
+                statistics["std"],
+                statistics["init_transition_prob"],
+            )
+            print(f" | > Data parameters loaded with value: {data_mean, data_std, init_transition_prob}")
+
+        trainer.config.flat_start_params["transition_p"] = (
+            init_transition_prob.item() if torch.is_tensor(init_transition_prob) else init_transition_prob
+        )
+        OverflowUtils.update_flat_start_transition(trainer.model, init_transition_prob)
+        trainer.model.update_mean_std(statistics)
+
+    @torch.inference_mode()
+    def _create_logs(self, batch, outputs, ap):  # pylint: disable=no-self-use, unused-argument
+        alignments, transition_vectors = outputs["alignments"], outputs["transition_vectors"]
+        means = torch.stack(outputs["means"], dim=1)
+
+        figures = {
+            "alignment": plot_alignment(alignments[0].exp(), title="Forward alignment", fig_size=(20, 20)),
+            "log_alignment": plot_alignment(
+                alignments[0].exp(), title="Forward log alignment", plot_log=True, fig_size=(20, 20)
+            ),
+            "transition_vectors": plot_alignment(transition_vectors[0], title="Transition vectors", fig_size=(20, 20)),
+            "mel_from_most_probable_state": plot_spectrogram(
+                get_spec_from_most_probable_state(alignments[0], means[0], self.decoder), fig_size=(12, 3)
+            ),
+            "mel_target": plot_spectrogram(batch["mel_input"][0], fig_size=(12, 3)),
+        }
+
+        # sample one item from the batch -1 will give the smalles item
+        print(" | > Synthesising audio from the model...")
+        inference_output = self.inference(
+            batch["text_input"][-1].unsqueeze(0), aux_input={"x_lengths": batch["text_lengths"][-1].unsqueeze(0)}
+        )
+        figures["synthesised"] = plot_spectrogram(inference_output["model_outputs"][0], fig_size=(12, 3))
+
+        states = [p[1] for p in inference_output["input_parameters"][0]]
+        transition_probability_synthesising = [p[2].cpu().numpy() for p in inference_output["output_parameters"][0]]
+
+        for i in range((len(transition_probability_synthesising) // 200) + 1):
+            start = i * 200
+            end = (i + 1) * 200
+            figures[f"synthesised_transition_probabilities/{i}"] = plot_transition_probabilities_to_numpy(
+                states[start:end], transition_probability_synthesising[start:end]
+            )
+
+        audio = ap.inv_melspectrogram(inference_output["model_outputs"][0].T.cpu().numpy())
+        return figures, {"audios": audio}
+
+    def train_log(
+        self, batch: dict, outputs: dict, logger: "Logger", assets: dict, steps: int
+    ):  # pylint: disable=unused-argument
+        """Log training progress."""
+        figures, audios = self._create_logs(batch, outputs, self.ap)
+        logger.train_figures(steps, figures)
+        logger.train_audios(steps, audios, self.ap.sample_rate)
+
+    def eval_log(
+        self, batch: Dict, outputs: Dict, logger: "Logger", assets: Dict, steps: int
+    ):  # pylint: disable=unused-argument
+        """Compute and log evaluation metrics."""
+        # Plot model parameters histograms
+        if isinstance(logger, TensorboardLogger):
+            # I don't know if any other loggers supports this
+            for tag, value in self.named_parameters():
+                tag = tag.replace(".", "/")
+                logger.writer.add_histogram(tag, value.data.cpu().numpy(), steps)
+
+        figures, audios = self._create_logs(batch, outputs, self.ap)
+        logger.eval_figures(steps, figures)
+        logger.eval_audios(steps, audios, self.ap.sample_rate)
+
+    def test_log(
+        self, outputs: dict, logger: "Logger", assets: dict, steps: int  # pylint: disable=unused-argument
+    ) -> None:
+        logger.test_audios(steps, outputs[1], self.ap.sample_rate)
+        logger.test_figures(steps, outputs[0])
+
+
+class NLLLoss(nn.Module):
+    """Negative log likelihood loss."""
+
+    def forward(self, log_prob: torch.Tensor) -> dict:  # pylint: disable=no-self-use
+        """Compute the loss.
+
+        Args:
+            logits (Tensor): [B, T, D]
+
+        Returns:
+            Tensor: [1]
+
+        """
+        return_dict = {}
+        return_dict["loss"] = -log_prob.mean()
+        return return_dict

TTS/tts/models/tacotron.py

@@ -0,0 +1,409 @@
+# coding: utf-8
+
+from typing import Dict, List, Tuple, Union
+
+import torch
+from torch import nn
+from torch.cuda.amp.autocast_mode import autocast
+from trainer.trainer_utils import get_optimizer, get_scheduler
+
+from TTS.tts.layers.tacotron.capacitron_layers import CapacitronVAE
+from TTS.tts.layers.tacotron.gst_layers import GST
+from TTS.tts.layers.tacotron.tacotron import Decoder, Encoder, PostCBHG
+from TTS.tts.models.base_tacotron import BaseTacotron
+from TTS.tts.utils.measures import alignment_diagonal_score
+from TTS.tts.utils.speakers import SpeakerManager
+from TTS.tts.utils.text.tokenizer import TTSTokenizer
+from TTS.tts.utils.visual import plot_alignment, plot_spectrogram
+from TTS.utils.capacitron_optimizer import CapacitronOptimizer
+
+
+class Tacotron(BaseTacotron):
+    """Tacotron as in https://arxiv.org/abs/1703.10135
+    It's an autoregressive encoder-attention-decoder-postnet architecture.
+    Check `TacotronConfig` for the arguments.
+
+    Args:
+        config (TacotronConfig): Configuration for the Tacotron model.
+        speaker_manager (SpeakerManager): Speaker manager to handle multi-speaker settings. Only use if the model is
+            a multi-speaker model. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        config: "TacotronConfig",
+        ap: "AudioProcessor" = None,
+        tokenizer: "TTSTokenizer" = None,
+        speaker_manager: SpeakerManager = None,
+    ):
+        super().__init__(config, ap, tokenizer, speaker_manager)
+
+        # pass all config fields to `self`
+        # for fewer code change
+        for key in config:
+            setattr(self, key, config[key])
+
+        # set speaker embedding channel size for determining `in_channels` for the connected layers.
+        # `init_multispeaker` needs to be called once more in training to initialize the speaker embedding layer based
+        # on the number of speakers infered from the dataset.
+        if self.use_speaker_embedding or self.use_d_vector_file:
+            self.init_multispeaker(config)
+            self.decoder_in_features += self.embedded_speaker_dim  # add speaker embedding dim
+
+        if self.use_gst:
+            self.decoder_in_features += self.gst.gst_embedding_dim
+
+        if self.use_capacitron_vae:
+            self.decoder_in_features += self.capacitron_vae.capacitron_VAE_embedding_dim
+
+        # embedding layer
+        self.embedding = nn.Embedding(self.num_chars, 256, padding_idx=0)
+        self.embedding.weight.data.normal_(0, 0.3)
+
+        # base model layers
+        self.encoder = Encoder(self.encoder_in_features)
+        self.decoder = Decoder(
+            self.decoder_in_features,
+            self.decoder_output_dim,
+            self.r,
+            self.memory_size,
+            self.attention_type,
+            self.windowing,
+            self.attention_norm,
+            self.prenet_type,
+            self.prenet_dropout,
+            self.use_forward_attn,
+            self.transition_agent,
+            self.forward_attn_mask,
+            self.location_attn,
+            self.attention_heads,
+            self.separate_stopnet,
+            self.max_decoder_steps,
+        )
+        self.postnet = PostCBHG(self.decoder_output_dim)
+        self.last_linear = nn.Linear(self.postnet.cbhg.gru_features * 2, self.out_channels)
+
+        # setup prenet dropout
+        self.decoder.prenet.dropout_at_inference = self.prenet_dropout_at_inference
+
+        # global style token layers
+        if self.gst and self.use_gst:
+            self.gst_layer = GST(
+                num_mel=self.decoder_output_dim,
+                num_heads=self.gst.gst_num_heads,
+                num_style_tokens=self.gst.gst_num_style_tokens,
+                gst_embedding_dim=self.gst.gst_embedding_dim,
+            )
+
+        # Capacitron layers
+        if self.capacitron_vae and self.use_capacitron_vae:
+            self.capacitron_vae_layer = CapacitronVAE(
+                num_mel=self.decoder_output_dim,
+                encoder_output_dim=self.encoder_in_features,
+                capacitron_VAE_embedding_dim=self.capacitron_vae.capacitron_VAE_embedding_dim,
+                speaker_embedding_dim=self.embedded_speaker_dim
+                if self.use_speaker_embedding and self.capacitron_vae.capacitron_use_speaker_embedding
+                else None,
+                text_summary_embedding_dim=self.capacitron_vae.capacitron_text_summary_embedding_dim
+                if self.capacitron_vae.capacitron_use_text_summary_embeddings
+                else None,
+            )
+
+        # backward pass decoder
+        if self.bidirectional_decoder:
+            self._init_backward_decoder()
+        # setup DDC
+        if self.double_decoder_consistency:
+            self.coarse_decoder = Decoder(
+                self.decoder_in_features,
+                self.decoder_output_dim,
+                self.ddc_r,
+                self.memory_size,
+                self.attention_type,
+                self.windowing,
+                self.attention_norm,
+                self.prenet_type,
+                self.prenet_dropout,
+                self.use_forward_attn,
+                self.transition_agent,
+                self.forward_attn_mask,
+                self.location_attn,
+                self.attention_heads,
+                self.separate_stopnet,
+                self.max_decoder_steps,
+            )
+
+    def forward(  # pylint: disable=dangerous-default-value
+        self, text, text_lengths, mel_specs=None, mel_lengths=None, aux_input={"speaker_ids": None, "d_vectors": None}
+    ):
+        """
+        Shapes:
+            text: [B, T_in]
+            text_lengths: [B]
+            mel_specs: [B, T_out, C]
+            mel_lengths: [B]
+            aux_input: 'speaker_ids': [B, 1] and  'd_vectors':[B, C]
+        """
+        aux_input = self._format_aux_input(aux_input)
+        outputs = {"alignments_backward": None, "decoder_outputs_backward": None}
+        inputs = self.embedding(text)
+        input_mask, output_mask = self.compute_masks(text_lengths, mel_lengths)
+        # B x T_in x encoder_in_features
+        encoder_outputs = self.encoder(inputs)
+        # sequence masking
+        encoder_outputs = encoder_outputs * input_mask.unsqueeze(2).expand_as(encoder_outputs)
+        # global style token
+        if self.gst and self.use_gst:
+            # B x gst_dim
+            encoder_outputs = self.compute_gst(encoder_outputs, mel_specs)
+        # speaker embedding
+        if self.use_speaker_embedding or self.use_d_vector_file:
+            if not self.use_d_vector_file:
+                # B x 1 x speaker_embed_dim
+                embedded_speakers = self.speaker_embedding(aux_input["speaker_ids"])[:, None]
+            else:
+                # B x 1 x speaker_embed_dim
+                embedded_speakers = torch.unsqueeze(aux_input["d_vectors"], 1)
+            encoder_outputs = self._concat_speaker_embedding(encoder_outputs, embedded_speakers)
+        # Capacitron
+        if self.capacitron_vae and self.use_capacitron_vae:
+            # B x capacitron_VAE_embedding_dim
+            encoder_outputs, *capacitron_vae_outputs = self.compute_capacitron_VAE_embedding(
+                encoder_outputs,
+                reference_mel_info=[mel_specs, mel_lengths],
+                text_info=[inputs, text_lengths]
+                if self.capacitron_vae.capacitron_use_text_summary_embeddings
+                else None,
+                speaker_embedding=embedded_speakers if self.capacitron_vae.capacitron_use_speaker_embedding else None,
+            )
+        else:
+            capacitron_vae_outputs = None
+        # decoder_outputs: B x decoder_in_features x T_out
+        # alignments: B x T_in x encoder_in_features
+        # stop_tokens: B x T_in
+        decoder_outputs, alignments, stop_tokens = self.decoder(encoder_outputs, mel_specs, input_mask)
+        # sequence masking
+        if output_mask is not None:
+            decoder_outputs = decoder_outputs * output_mask.unsqueeze(1).expand_as(decoder_outputs)
+        # B x T_out x decoder_in_features
+        postnet_outputs = self.postnet(decoder_outputs)
+        # sequence masking
+        if output_mask is not None:
+            postnet_outputs = postnet_outputs * output_mask.unsqueeze(2).expand_as(postnet_outputs)
+        # B x T_out x posnet_dim
+        postnet_outputs = self.last_linear(postnet_outputs)
+        # B x T_out x decoder_in_features
+        decoder_outputs = decoder_outputs.transpose(1, 2).contiguous()
+        if self.bidirectional_decoder:
+            decoder_outputs_backward, alignments_backward = self._backward_pass(mel_specs, encoder_outputs, input_mask)
+            outputs["alignments_backward"] = alignments_backward
+            outputs["decoder_outputs_backward"] = decoder_outputs_backward
+        if self.double_decoder_consistency:
+            decoder_outputs_backward, alignments_backward = self._coarse_decoder_pass(
+                mel_specs, encoder_outputs, alignments, input_mask
+            )
+            outputs["alignments_backward"] = alignments_backward
+            outputs["decoder_outputs_backward"] = decoder_outputs_backward
+        outputs.update(
+            {
+                "model_outputs": postnet_outputs,
+                "decoder_outputs": decoder_outputs,
+                "alignments": alignments,
+                "stop_tokens": stop_tokens,
+                "capacitron_vae_outputs": capacitron_vae_outputs,
+            }
+        )
+        return outputs
+
+    @torch.no_grad()
+    def inference(self, text_input, aux_input=None):
+        aux_input = self._format_aux_input(aux_input)
+        inputs = self.embedding(text_input)
+        encoder_outputs = self.encoder(inputs)
+        if self.gst and self.use_gst:
+            # B x gst_dim
+            encoder_outputs = self.compute_gst(encoder_outputs, aux_input["style_mel"], aux_input["d_vectors"])
+        if self.capacitron_vae and self.use_capacitron_vae:
+            if aux_input["style_text"] is not None:
+                style_text_embedding = self.embedding(aux_input["style_text"])
+                style_text_length = torch.tensor([style_text_embedding.size(1)], dtype=torch.int64).to(
+                    encoder_outputs.device
+                )  # pylint: disable=not-callable
+            reference_mel_length = (
+                torch.tensor([aux_input["style_mel"].size(1)], dtype=torch.int64).to(encoder_outputs.device)
+                if aux_input["style_mel"] is not None
+                else None
+            )  # pylint: disable=not-callable
+            # B x capacitron_VAE_embedding_dim
+            encoder_outputs, *_ = self.compute_capacitron_VAE_embedding(
+                encoder_outputs,
+                reference_mel_info=[aux_input["style_mel"], reference_mel_length]
+                if aux_input["style_mel"] is not None
+                else None,
+                text_info=[style_text_embedding, style_text_length] if aux_input["style_text"] is not None else None,
+                speaker_embedding=aux_input["d_vectors"]
+                if self.capacitron_vae.capacitron_use_speaker_embedding
+                else None,
+            )
+        if self.num_speakers > 1:
+            if not self.use_d_vector_file:
+                # B x 1 x speaker_embed_dim
+                embedded_speakers = self.speaker_embedding(aux_input["speaker_ids"])
+                # reshape embedded_speakers
+                if embedded_speakers.ndim == 1:
+                    embedded_speakers = embedded_speakers[None, None, :]
+                elif embedded_speakers.ndim == 2:
+                    embedded_speakers = embedded_speakers[None, :]
+            else:
+                # B x 1 x speaker_embed_dim
+                embedded_speakers = torch.unsqueeze(aux_input["d_vectors"], 1)
+            encoder_outputs = self._concat_speaker_embedding(encoder_outputs, embedded_speakers)
+        decoder_outputs, alignments, stop_tokens = self.decoder.inference(encoder_outputs)
+        postnet_outputs = self.postnet(decoder_outputs)
+        postnet_outputs = self.last_linear(postnet_outputs)
+        decoder_outputs = decoder_outputs.transpose(1, 2)
+        outputs = {
+            "model_outputs": postnet_outputs,
+            "decoder_outputs": decoder_outputs,
+            "alignments": alignments,
+            "stop_tokens": stop_tokens,
+        }
+        return outputs
+
+    def before_backward_pass(self, loss_dict, optimizer) -> None:
+        # Extracting custom training specific operations for capacitron
+        # from the trainer
+        if self.use_capacitron_vae:
+            loss_dict["capacitron_vae_beta_loss"].backward()
+            optimizer.first_step()
+
+    def train_step(self, batch: Dict, criterion: torch.nn.Module) -> Tuple[Dict, Dict]:
+        """Perform a single training step by fetching the right set of samples from the batch.
+
+        Args:
+            batch ([Dict]): A dictionary of input tensors.
+            criterion ([torch.nn.Module]): Callable criterion to compute model loss.
+        """
+        text_input = batch["text_input"]
+        text_lengths = batch["text_lengths"]
+        mel_input = batch["mel_input"]
+        mel_lengths = batch["mel_lengths"]
+        linear_input = batch["linear_input"]
+        stop_targets = batch["stop_targets"]
+        stop_target_lengths = batch["stop_target_lengths"]
+        speaker_ids = batch["speaker_ids"]
+        d_vectors = batch["d_vectors"]
+
+        aux_input = {"speaker_ids": speaker_ids, "d_vectors": d_vectors}
+        outputs = self.forward(text_input, text_lengths, mel_input, mel_lengths, aux_input)
+
+        # set the [alignment] lengths wrt reduction factor for guided attention
+        if mel_lengths.max() % self.decoder.r != 0:
+            alignment_lengths = (
+                mel_lengths + (self.decoder.r - (mel_lengths.max() % self.decoder.r))
+            ) // self.decoder.r
+        else:
+            alignment_lengths = mel_lengths // self.decoder.r
+
+        # compute loss
+        with autocast(enabled=False):  # use float32 for the criterion
+            loss_dict = criterion(
+                outputs["model_outputs"].float(),
+                outputs["decoder_outputs"].float(),
+                mel_input.float(),
+                linear_input.float(),
+                outputs["stop_tokens"].float(),
+                stop_targets.float(),
+                stop_target_lengths,
+                outputs["capacitron_vae_outputs"] if self.capacitron_vae else None,
+                mel_lengths,
+                None if outputs["decoder_outputs_backward"] is None else outputs["decoder_outputs_backward"].float(),
+                outputs["alignments"].float(),
+                alignment_lengths,
+                None if outputs["alignments_backward"] is None else outputs["alignments_backward"].float(),
+                text_lengths,
+            )
+
+        # compute alignment error (the lower the better )
+        align_error = 1 - alignment_diagonal_score(outputs["alignments"])
+        loss_dict["align_error"] = align_error
+        return outputs, loss_dict
+
+    def get_optimizer(self) -> List:
+        if self.use_capacitron_vae:
+            return CapacitronOptimizer(self.config, self.named_parameters())
+        return get_optimizer(self.config.optimizer, self.config.optimizer_params, self.config.lr, self)
+
+    def get_scheduler(self, optimizer: object):
+        opt = optimizer.primary_optimizer if self.use_capacitron_vae else optimizer
+        return get_scheduler(self.config.lr_scheduler, self.config.lr_scheduler_params, opt)
+
+    def before_gradient_clipping(self):
+        if self.use_capacitron_vae:
+            # Capacitron model specific gradient clipping
+            model_params_to_clip = []
+            for name, param in self.named_parameters():
+                if param.requires_grad:
+                    if name != "capacitron_vae_layer.beta":
+                        model_params_to_clip.append(param)
+            torch.nn.utils.clip_grad_norm_(model_params_to_clip, self.capacitron_vae.capacitron_grad_clip)
+
+    def _create_logs(self, batch, outputs, ap):
+        postnet_outputs = outputs["model_outputs"]
+        decoder_outputs = outputs["decoder_outputs"]
+        alignments = outputs["alignments"]
+        alignments_backward = outputs["alignments_backward"]
+        mel_input = batch["mel_input"]
+        linear_input = batch["linear_input"]
+
+        pred_linear_spec = postnet_outputs[0].data.cpu().numpy()
+        pred_mel_spec = decoder_outputs[0].data.cpu().numpy()
+        gt_linear_spec = linear_input[0].data.cpu().numpy()
+        gt_mel_spec = mel_input[0].data.cpu().numpy()
+        align_img = alignments[0].data.cpu().numpy()
+
+        figures = {
+            "pred_linear_spec": plot_spectrogram(pred_linear_spec, ap, output_fig=False),
+            "real_linear_spec": plot_spectrogram(gt_linear_spec, ap, output_fig=False),
+            "pred_mel_spec": plot_spectrogram(pred_mel_spec, ap, output_fig=False),
+            "real_mel_spec": plot_spectrogram(gt_mel_spec, ap, output_fig=False),
+            "alignment": plot_alignment(align_img, output_fig=False),
+        }
+
+        if self.bidirectional_decoder or self.double_decoder_consistency:
+            figures["alignment_backward"] = plot_alignment(alignments_backward[0].data.cpu().numpy(), output_fig=False)
+
+        # Sample audio
+        audio = ap.inv_spectrogram(pred_linear_spec.T)
+        return figures, {"audio": audio}
+
+    def train_log(
+        self, batch: dict, outputs: dict, logger: "Logger", assets: dict, steps: int
+    ) -> None:  # pylint: disable=no-self-use
+        figures, audios = self._create_logs(batch, outputs, self.ap)
+        logger.train_figures(steps, figures)
+        logger.train_audios(steps, audios, self.ap.sample_rate)
+
+    def eval_step(self, batch: dict, criterion: nn.Module):
+        return self.train_step(batch, criterion)
+
+    def eval_log(self, batch: dict, outputs: dict, logger: "Logger", assets: dict, steps: int) -> None:
+        figures, audios = self._create_logs(batch, outputs, self.ap)
+        logger.eval_figures(steps, figures)
+        logger.eval_audios(steps, audios, self.ap.sample_rate)
+
+    @staticmethod
+    def init_from_config(config: "TacotronConfig", samples: Union[List[List], List[Dict]] = None):
+        """Initiate model from config
+
+        Args:
+            config (TacotronConfig): Model config.
+            samples (Union[List[List], List[Dict]]): Training samples to parse speaker ids for training.
+                Defaults to None.
+        """
+        from TTS.utils.audio import AudioProcessor
+
+        ap = AudioProcessor.init_from_config(config)
+        tokenizer, new_config = TTSTokenizer.init_from_config(config)
+        speaker_manager = SpeakerManager.init_from_config(config, samples)
+        return Tacotron(new_config, ap, tokenizer, speaker_manager)

TTS/tts/models/tacotron2.py

@@ -0,0 +1,433 @@
+# coding: utf-8
+
+from typing import Dict, List, Union
+
+import torch
+from torch import nn
+from torch.cuda.amp.autocast_mode import autocast
+from trainer.trainer_utils import get_optimizer, get_scheduler
+
+from TTS.tts.layers.tacotron.capacitron_layers import CapacitronVAE
+from TTS.tts.layers.tacotron.gst_layers import GST
+from TTS.tts.layers.tacotron.tacotron2 import Decoder, Encoder, Postnet
+from TTS.tts.models.base_tacotron import BaseTacotron
+from TTS.tts.utils.measures import alignment_diagonal_score
+from TTS.tts.utils.speakers import SpeakerManager
+from TTS.tts.utils.text.tokenizer import TTSTokenizer
+from TTS.tts.utils.visual import plot_alignment, plot_spectrogram
+from TTS.utils.capacitron_optimizer import CapacitronOptimizer
+
+
+class Tacotron2(BaseTacotron):
+    """Tacotron2 model implementation inherited from :class:`TTS.tts.models.base_tacotron.BaseTacotron`.
+
+    Paper::
+        https://arxiv.org/abs/1712.05884
+
+    Paper abstract::
+        This paper describes Tacotron 2, a neural network architecture for speech synthesis directly from text.
+        The system is composed of a recurrent sequence-to-sequence feature prediction network that maps character
+        embeddings to mel-scale spectrograms, followed by a modified WaveNet model acting as a vocoder to synthesize
+        timedomain waveforms from those spectrograms. Our model achieves a mean opinion score (MOS) of 4.53 comparable
+        to a MOS of 4.58 for professionally recorded speech. To validate our design choices, we present ablation
+        studies of key components of our system and evaluate the impact of using mel spectrograms as the input to
+        WaveNet instead of linguistic, duration, and F0 features. We further demonstrate that using a compact acoustic
+        intermediate representation enables significant simplification of the WaveNet architecture.
+
+    Check :class:`TTS.tts.configs.tacotron2_config.Tacotron2Config` for model arguments.
+
+    Args:
+        config (TacotronConfig):
+            Configuration for the Tacotron2 model.
+        speaker_manager (SpeakerManager):
+            Speaker manager for multi-speaker training. Uuse only for multi-speaker training. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        config: "Tacotron2Config",
+        ap: "AudioProcessor" = None,
+        tokenizer: "TTSTokenizer" = None,
+        speaker_manager: SpeakerManager = None,
+    ):
+        super().__init__(config, ap, tokenizer, speaker_manager)
+
+        self.decoder_output_dim = config.out_channels
+
+        # pass all config fields to `self`
+        # for fewer code change
+        for key in config:
+            setattr(self, key, config[key])
+
+        # init multi-speaker layers
+        if self.use_speaker_embedding or self.use_d_vector_file:
+            self.init_multispeaker(config)
+            self.decoder_in_features += self.embedded_speaker_dim  # add speaker embedding dim
+
+        if self.use_gst:
+            self.decoder_in_features += self.gst.gst_embedding_dim
+
+        if self.use_capacitron_vae:
+            self.decoder_in_features += self.capacitron_vae.capacitron_VAE_embedding_dim
+
+        # embedding layer
+        self.embedding = nn.Embedding(self.num_chars, 512, padding_idx=0)
+
+        # base model layers
+        self.encoder = Encoder(self.encoder_in_features)
+
+        self.decoder = Decoder(
+            self.decoder_in_features,
+            self.decoder_output_dim,
+            self.r,
+            self.attention_type,
+            self.attention_win,
+            self.attention_norm,
+            self.prenet_type,
+            self.prenet_dropout,
+            self.use_forward_attn,
+            self.transition_agent,
+            self.forward_attn_mask,
+            self.location_attn,
+            self.attention_heads,
+            self.separate_stopnet,
+            self.max_decoder_steps,
+        )
+        self.postnet = Postnet(self.out_channels)
+
+        # setup prenet dropout
+        self.decoder.prenet.dropout_at_inference = self.prenet_dropout_at_inference
+
+        # global style token layers
+        if self.gst and self.use_gst:
+            self.gst_layer = GST(
+                num_mel=self.decoder_output_dim,
+                num_heads=self.gst.gst_num_heads,
+                num_style_tokens=self.gst.gst_num_style_tokens,
+                gst_embedding_dim=self.gst.gst_embedding_dim,
+            )
+
+        # Capacitron VAE Layers
+        if self.capacitron_vae and self.use_capacitron_vae:
+            self.capacitron_vae_layer = CapacitronVAE(
+                num_mel=self.decoder_output_dim,
+                encoder_output_dim=self.encoder_in_features,
+                capacitron_VAE_embedding_dim=self.capacitron_vae.capacitron_VAE_embedding_dim,
+                speaker_embedding_dim=self.embedded_speaker_dim
+                if self.capacitron_vae.capacitron_use_speaker_embedding
+                else None,
+                text_summary_embedding_dim=self.capacitron_vae.capacitron_text_summary_embedding_dim
+                if self.capacitron_vae.capacitron_use_text_summary_embeddings
+                else None,
+            )
+
+        # backward pass decoder
+        if self.bidirectional_decoder:
+            self._init_backward_decoder()
+        # setup DDC
+        if self.double_decoder_consistency:
+            self.coarse_decoder = Decoder(
+                self.decoder_in_features,
+                self.decoder_output_dim,
+                self.ddc_r,
+                self.attention_type,
+                self.attention_win,
+                self.attention_norm,
+                self.prenet_type,
+                self.prenet_dropout,
+                self.use_forward_attn,
+                self.transition_agent,
+                self.forward_attn_mask,
+                self.location_attn,
+                self.attention_heads,
+                self.separate_stopnet,
+                self.max_decoder_steps,
+            )
+
+    @staticmethod
+    def shape_outputs(mel_outputs, mel_outputs_postnet, alignments):
+        """Final reshape of the model output tensors."""
+        mel_outputs = mel_outputs.transpose(1, 2)
+        mel_outputs_postnet = mel_outputs_postnet.transpose(1, 2)
+        return mel_outputs, mel_outputs_postnet, alignments
+
+    def forward(  # pylint: disable=dangerous-default-value
+        self, text, text_lengths, mel_specs=None, mel_lengths=None, aux_input={"speaker_ids": None, "d_vectors": None}
+    ):
+        """Forward pass for training with Teacher Forcing.
+
+        Shapes:
+            text: :math:`[B, T_in]`
+            text_lengths: :math:`[B]`
+            mel_specs: :math:`[B, T_out, C]`
+            mel_lengths: :math:`[B]`
+            aux_input: 'speaker_ids': :math:`[B, 1]` and  'd_vectors': :math:`[B, C]`
+        """
+        aux_input = self._format_aux_input(aux_input)
+        outputs = {"alignments_backward": None, "decoder_outputs_backward": None}
+        # compute mask for padding
+        # B x T_in_max (boolean)
+        input_mask, output_mask = self.compute_masks(text_lengths, mel_lengths)
+        # B x D_embed x T_in_max
+        embedded_inputs = self.embedding(text).transpose(1, 2)
+        # B x T_in_max x D_en
+        encoder_outputs = self.encoder(embedded_inputs, text_lengths)
+        if self.gst and self.use_gst:
+            # B x gst_dim
+            encoder_outputs = self.compute_gst(encoder_outputs, mel_specs)
+
+        if self.use_speaker_embedding or self.use_d_vector_file:
+            if not self.use_d_vector_file:
+                # B x 1 x speaker_embed_dim
+                embedded_speakers = self.speaker_embedding(aux_input["speaker_ids"])[:, None]
+            else:
+                # B x 1 x speaker_embed_dim
+                embedded_speakers = torch.unsqueeze(aux_input["d_vectors"], 1)
+            encoder_outputs = self._concat_speaker_embedding(encoder_outputs, embedded_speakers)
+
+        # capacitron
+        if self.capacitron_vae and self.use_capacitron_vae:
+            # B x capacitron_VAE_embedding_dim
+            encoder_outputs, *capacitron_vae_outputs = self.compute_capacitron_VAE_embedding(
+                encoder_outputs,
+                reference_mel_info=[mel_specs, mel_lengths],
+                text_info=[embedded_inputs.transpose(1, 2), text_lengths]
+                if self.capacitron_vae.capacitron_use_text_summary_embeddings
+                else None,
+                speaker_embedding=embedded_speakers if self.capacitron_vae.capacitron_use_speaker_embedding else None,
+            )
+        else:
+            capacitron_vae_outputs = None
+
+        encoder_outputs = encoder_outputs * input_mask.unsqueeze(2).expand_as(encoder_outputs)
+
+        # B x mel_dim x T_out -- B x T_out//r x T_in -- B x T_out//r
+        decoder_outputs, alignments, stop_tokens = self.decoder(encoder_outputs, mel_specs, input_mask)
+        # sequence masking
+        if mel_lengths is not None:
+            decoder_outputs = decoder_outputs * output_mask.unsqueeze(1).expand_as(decoder_outputs)
+        # B x mel_dim x T_out
+        postnet_outputs = self.postnet(decoder_outputs)
+        postnet_outputs = decoder_outputs + postnet_outputs
+        # sequence masking
+        if output_mask is not None:
+            postnet_outputs = postnet_outputs * output_mask.unsqueeze(1).expand_as(postnet_outputs)
+        # B x T_out x mel_dim -- B x T_out x mel_dim -- B x T_out//r x T_in
+        decoder_outputs, postnet_outputs, alignments = self.shape_outputs(decoder_outputs, postnet_outputs, alignments)
+        if self.bidirectional_decoder:
+            decoder_outputs_backward, alignments_backward = self._backward_pass(mel_specs, encoder_outputs, input_mask)
+            outputs["alignments_backward"] = alignments_backward
+            outputs["decoder_outputs_backward"] = decoder_outputs_backward
+        if self.double_decoder_consistency:
+            decoder_outputs_backward, alignments_backward = self._coarse_decoder_pass(
+                mel_specs, encoder_outputs, alignments, input_mask
+            )
+            outputs["alignments_backward"] = alignments_backward
+            outputs["decoder_outputs_backward"] = decoder_outputs_backward
+        outputs.update(
+            {
+                "model_outputs": postnet_outputs,
+                "decoder_outputs": decoder_outputs,
+                "alignments": alignments,
+                "stop_tokens": stop_tokens,
+                "capacitron_vae_outputs": capacitron_vae_outputs,
+            }
+        )
+        return outputs
+
+    @torch.no_grad()
+    def inference(self, text, aux_input=None):
+        """Forward pass for inference with no Teacher-Forcing.
+
+        Shapes:
+           text: :math:`[B, T_in]`
+           text_lengths: :math:`[B]`
+        """
+        aux_input = self._format_aux_input(aux_input)
+        embedded_inputs = self.embedding(text).transpose(1, 2)
+        encoder_outputs = self.encoder.inference(embedded_inputs)
+
+        if self.gst and self.use_gst:
+            # B x gst_dim
+            encoder_outputs = self.compute_gst(encoder_outputs, aux_input["style_mel"], aux_input["d_vectors"])
+
+        if self.capacitron_vae and self.use_capacitron_vae:
+            if aux_input["style_text"] is not None:
+                style_text_embedding = self.embedding(aux_input["style_text"])
+                style_text_length = torch.tensor([style_text_embedding.size(1)], dtype=torch.int64).to(
+                    encoder_outputs.device
+                )  # pylint: disable=not-callable
+            reference_mel_length = (
+                torch.tensor([aux_input["style_mel"].size(1)], dtype=torch.int64).to(encoder_outputs.device)
+                if aux_input["style_mel"] is not None
+                else None
+            )  # pylint: disable=not-callable
+            # B x capacitron_VAE_embedding_dim
+            encoder_outputs, *_ = self.compute_capacitron_VAE_embedding(
+                encoder_outputs,
+                reference_mel_info=[aux_input["style_mel"], reference_mel_length]
+                if aux_input["style_mel"] is not None
+                else None,
+                text_info=[style_text_embedding, style_text_length] if aux_input["style_text"] is not None else None,
+                speaker_embedding=aux_input["d_vectors"]
+                if self.capacitron_vae.capacitron_use_speaker_embedding
+                else None,
+            )
+
+        if self.num_speakers > 1:
+            if not self.use_d_vector_file:
+                embedded_speakers = self.speaker_embedding(aux_input["speaker_ids"])[None]
+                # reshape embedded_speakers
+                if embedded_speakers.ndim == 1:
+                    embedded_speakers = embedded_speakers[None, None, :]
+                elif embedded_speakers.ndim == 2:
+                    embedded_speakers = embedded_speakers[None, :]
+            else:
+                embedded_speakers = aux_input["d_vectors"]
+
+            encoder_outputs = self._concat_speaker_embedding(encoder_outputs, embedded_speakers)
+
+        decoder_outputs, alignments, stop_tokens = self.decoder.inference(encoder_outputs)
+        postnet_outputs = self.postnet(decoder_outputs)
+        postnet_outputs = decoder_outputs + postnet_outputs
+        decoder_outputs, postnet_outputs, alignments = self.shape_outputs(decoder_outputs, postnet_outputs, alignments)
+        outputs = {
+            "model_outputs": postnet_outputs,
+            "decoder_outputs": decoder_outputs,
+            "alignments": alignments,
+            "stop_tokens": stop_tokens,
+        }
+        return outputs
+
+    def before_backward_pass(self, loss_dict, optimizer) -> None:
+        # Extracting custom training specific operations for capacitron
+        # from the trainer
+        if self.use_capacitron_vae:
+            loss_dict["capacitron_vae_beta_loss"].backward()
+            optimizer.first_step()
+
+    def train_step(self, batch: Dict, criterion: torch.nn.Module):
+        """A single training step. Forward pass and loss computation.
+
+        Args:
+            batch ([Dict]): A dictionary of input tensors.
+            criterion ([type]): Callable criterion to compute model loss.
+        """
+        text_input = batch["text_input"]
+        text_lengths = batch["text_lengths"]
+        mel_input = batch["mel_input"]
+        mel_lengths = batch["mel_lengths"]
+        stop_targets = batch["stop_targets"]
+        stop_target_lengths = batch["stop_target_lengths"]
+        speaker_ids = batch["speaker_ids"]
+        d_vectors = batch["d_vectors"]
+
+        aux_input = {"speaker_ids": speaker_ids, "d_vectors": d_vectors}
+        outputs = self.forward(text_input, text_lengths, mel_input, mel_lengths, aux_input)
+
+        # set the [alignment] lengths wrt reduction factor for guided attention
+        if mel_lengths.max() % self.decoder.r != 0:
+            alignment_lengths = (
+                mel_lengths + (self.decoder.r - (mel_lengths.max() % self.decoder.r))
+            ) // self.decoder.r
+        else:
+            alignment_lengths = mel_lengths // self.decoder.r
+
+        # compute loss
+        with autocast(enabled=False):  # use float32 for the criterion
+            loss_dict = criterion(
+                outputs["model_outputs"].float(),
+                outputs["decoder_outputs"].float(),
+                mel_input.float(),
+                None,
+                outputs["stop_tokens"].float(),
+                stop_targets.float(),
+                stop_target_lengths,
+                outputs["capacitron_vae_outputs"] if self.capacitron_vae else None,
+                mel_lengths,
+                None if outputs["decoder_outputs_backward"] is None else outputs["decoder_outputs_backward"].float(),
+                outputs["alignments"].float(),
+                alignment_lengths,
+                None if outputs["alignments_backward"] is None else outputs["alignments_backward"].float(),
+                text_lengths,
+            )
+
+        # compute alignment error (the lower the better )
+        align_error = 1 - alignment_diagonal_score(outputs["alignments"])
+        loss_dict["align_error"] = align_error
+        return outputs, loss_dict
+
+    def get_optimizer(self) -> List:
+        if self.use_capacitron_vae:
+            return CapacitronOptimizer(self.config, self.named_parameters())
+        return get_optimizer(self.config.optimizer, self.config.optimizer_params, self.config.lr, self)
+
+    def get_scheduler(self, optimizer: object):
+        opt = optimizer.primary_optimizer if self.use_capacitron_vae else optimizer
+        return get_scheduler(self.config.lr_scheduler, self.config.lr_scheduler_params, opt)
+
+    def before_gradient_clipping(self):
+        if self.use_capacitron_vae:
+            # Capacitron model specific gradient clipping
+            model_params_to_clip = []
+            for name, param in self.named_parameters():
+                if param.requires_grad:
+                    if name != "capacitron_vae_layer.beta":
+                        model_params_to_clip.append(param)
+            torch.nn.utils.clip_grad_norm_(model_params_to_clip, self.capacitron_vae.capacitron_grad_clip)
+
+    def _create_logs(self, batch, outputs, ap):
+        """Create dashboard log information."""
+        postnet_outputs = outputs["model_outputs"]
+        alignments = outputs["alignments"]
+        alignments_backward = outputs["alignments_backward"]
+        mel_input = batch["mel_input"]
+
+        pred_spec = postnet_outputs[0].data.cpu().numpy()
+        gt_spec = mel_input[0].data.cpu().numpy()
+        align_img = alignments[0].data.cpu().numpy()
+
+        figures = {
+            "prediction": plot_spectrogram(pred_spec, ap, output_fig=False),
+            "ground_truth": plot_spectrogram(gt_spec, ap, output_fig=False),
+            "alignment": plot_alignment(align_img, output_fig=False),
+        }
+
+        if self.bidirectional_decoder or self.double_decoder_consistency:
+            figures["alignment_backward"] = plot_alignment(alignments_backward[0].data.cpu().numpy(), output_fig=False)
+
+        # Sample audio
+        audio = ap.inv_melspectrogram(pred_spec.T)
+        return figures, {"audio": audio}
+
+    def train_log(
+        self, batch: dict, outputs: dict, logger: "Logger", assets: dict, steps: int
+    ) -> None:  # pylint: disable=no-self-use
+        """Log training progress."""
+        figures, audios = self._create_logs(batch, outputs, self.ap)
+        logger.train_figures(steps, figures)
+        logger.train_audios(steps, audios, self.ap.sample_rate)
+
+    def eval_step(self, batch: dict, criterion: nn.Module):
+        return self.train_step(batch, criterion)
+
+    def eval_log(self, batch: dict, outputs: dict, logger: "Logger", assets: dict, steps: int) -> None:
+        figures, audios = self._create_logs(batch, outputs, self.ap)
+        logger.eval_figures(steps, figures)
+        logger.eval_audios(steps, audios, self.ap.sample_rate)
+
+    @staticmethod
+    def init_from_config(config: "Tacotron2Config", samples: Union[List[List], List[Dict]] = None):
+        """Initiate model from config
+
+        Args:
+            config (Tacotron2Config): Model config.
+            samples (Union[List[List], List[Dict]]): Training samples to parse speaker ids for training.
+                Defaults to None.
+        """
+        from TTS.utils.audio import AudioProcessor
+
+        ap = AudioProcessor.init_from_config(config)
+        tokenizer, new_config = TTSTokenizer.init_from_config(config)
+        speaker_manager = SpeakerManager.init_from_config(new_config, samples)
+        return Tacotron2(new_config, ap, tokenizer, speaker_manager)

TTS/tts/models/tortoise.py

@@ -0,0 +1,911 @@
+import os
+import random
+from contextlib import contextmanager
+from dataclasses import dataclass
+from time import time
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+from coqpit import Coqpit
+from tqdm import tqdm
+
+from TTS.tts.layers.tortoise.arch_utils import TorchMelSpectrogram
+from TTS.tts.layers.tortoise.audio_utils import denormalize_tacotron_mel, load_voice, wav_to_univnet_mel
+from TTS.tts.layers.tortoise.autoregressive import UnifiedVoice
+from TTS.tts.layers.tortoise.classifier import AudioMiniEncoderWithClassifierHead
+from TTS.tts.layers.tortoise.clvp import CLVP
+from TTS.tts.layers.tortoise.diffusion import SpacedDiffusion, get_named_beta_schedule, space_timesteps
+from TTS.tts.layers.tortoise.diffusion_decoder import DiffusionTts
+from TTS.tts.layers.tortoise.random_latent_generator import RandomLatentConverter
+from TTS.tts.layers.tortoise.tokenizer import VoiceBpeTokenizer
+from TTS.tts.layers.tortoise.vocoder import VocConf, VocType
+from TTS.tts.layers.tortoise.wav2vec_alignment import Wav2VecAlignment
+from TTS.tts.models.base_tts import BaseTTS
+
+
+def pad_or_truncate(t, length):
+    """
+    Utility function for forcing <t> to have the specified sequence length, whether by clipping it or padding it with 0s.
+    """
+    tp = t[..., :length]
+    if t.shape[-1] == length:
+        tp = t
+    elif t.shape[-1] < length:
+        tp = F.pad(t, (0, length - t.shape[-1]))
+    return tp
+
+
+def deterministic_state(seed=None):
+    """
+    Sets the random seeds that tortoise uses to the current time() and returns that seed so results can be
+    reproduced.
+    """
+    seed = int(time()) if seed is None else seed
+    torch.manual_seed(seed)
+    random.seed(seed)
+    # Can't currently set this because of CUBLAS. TODO: potentially enable it if necessary.
+    # torch.use_deterministic_algorithms(True)
+
+    return seed
+
+
+def load_discrete_vocoder_diffuser(
+    trained_diffusion_steps=4000,
+    desired_diffusion_steps=200,
+    cond_free=True,
+    cond_free_k=1,
+    sampler="ddim",
+):
+    """
+    Helper function to load a GaussianDiffusion instance configured for use as a vocoder.
+    """
+    return SpacedDiffusion(
+        use_timesteps=space_timesteps(trained_diffusion_steps, [desired_diffusion_steps]),
+        model_mean_type="epsilon",
+        model_var_type="learned_range",
+        loss_type="mse",
+        betas=get_named_beta_schedule("linear", trained_diffusion_steps),
+        conditioning_free=cond_free,
+        conditioning_free_k=cond_free_k,
+        sampler=sampler,
+    )
+
+
+def format_conditioning(clip, cond_length=132300, device="cuda", **kwargs):
+    """
+    Converts the given conditioning signal to a MEL spectrogram and clips it as expected by the models.
+    """
+    gap = clip.shape[-1] - cond_length
+    if gap < 0:
+        clip = F.pad(clip, pad=(0, abs(gap)))
+    elif gap > 0:
+        rand_start = random.randint(0, gap)
+        clip = clip[:, rand_start : rand_start + cond_length]
+    mel_clip = TorchMelSpectrogram(**kwargs)(clip.unsqueeze(0)).squeeze(0)
+    return mel_clip.unsqueeze(0).to(device)
+
+
+def fix_autoregressive_output(codes, stop_token, complain=True):
+    """
+    This function performs some padding on coded audio that fixes a mismatch issue between what the diffusion model was
+    trained on and what the autoregressive code generator creates (which has no padding or end).
+    This is highly specific to the DVAE being used, so this particular coding will not necessarily work if used with
+    a different DVAE. This can be inferred by feeding a audio clip padded with lots of zeros on the end through the DVAE
+    and copying out the last few codes.
+
+    Failing to do this padding will produce speech with a harsh end that sounds like "BLAH" or similar.
+    """
+    # Strip off the autoregressive stop token and add padding.
+    stop_token_indices = (codes == stop_token).nonzero()
+    if len(stop_token_indices) == 0:
+        if complain:
+            print(
+                "No stop tokens found in one of the generated voice clips. This typically means the spoken audio is "
+                "too long. In some cases, the output will still be good, though. Listen to it and if it is missing words, "
+                "try breaking up your input text."
+            )
+        return codes
+    codes[stop_token_indices] = 83
+    stm = stop_token_indices.min().item()
+    codes[stm:] = 83
+    if stm - 3 < codes.shape[0]:
+        codes[-3] = 45
+        codes[-2] = 45
+        codes[-1] = 248
+    return codes
+
+
+def do_spectrogram_diffusion(
+    diffusion_model,
+    diffuser,
+    latents,
+    conditioning_latents,
+    temperature=1,
+    verbose=True,
+):
+    """
+    Uses the specified diffusion model to convert discrete codes into a spectrogram.
+    """
+    with torch.no_grad():
+        output_seq_len = (
+            latents.shape[1] * 4 * 24000 // 22050
+        )  # This diffusion model converts from 22kHz spectrogram codes to a 24kHz spectrogram signal.
+        output_shape = (latents.shape[0], 100, output_seq_len)
+        precomputed_embeddings = diffusion_model.timestep_independent(
+            latents, conditioning_latents, output_seq_len, False
+        )
+
+        noise = torch.randn(output_shape, device=latents.device) * temperature
+        mel = diffuser.sample_loop(
+            diffusion_model,
+            output_shape,
+            noise=noise,
+            model_kwargs={"precomputed_aligned_embeddings": precomputed_embeddings},
+            progress=verbose,
+        )
+        return denormalize_tacotron_mel(mel)[:, :, :output_seq_len]
+
+
+def classify_audio_clip(clip, model_dir):
+    """
+    Returns whether or not Tortoises' classifier thinks the given clip came from Tortoise.
+    :param clip: torch tensor containing audio waveform data (get it from load_audio)
+    :return: True if the clip was classified as coming from Tortoise and false if it was classified as real.
+    """
+    classifier = AudioMiniEncoderWithClassifierHead(
+        2,
+        spec_dim=1,
+        embedding_dim=512,
+        depth=5,
+        downsample_factor=4,
+        resnet_blocks=2,
+        attn_blocks=4,
+        num_attn_heads=4,
+        base_channels=32,
+        dropout=0,
+        kernel_size=5,
+        distribute_zero_label=False,
+    )
+    classifier.load_state_dict(torch.load(os.path.join(model_dir, "classifier.pth"), map_location=torch.device("cpu")))
+    clip = clip.cpu().unsqueeze(0)
+    results = F.softmax(classifier(clip), dim=-1)
+    return results[0][0]
+
+
+def pick_best_batch_size_for_gpu():
+    """
+    Tries to pick a batch size that will fit in your GPU. These sizes aren't guaranteed to work, but they should give
+    you a good shot.
+    """
+    if torch.cuda.is_available():
+        _, available = torch.cuda.mem_get_info()
+        availableGb = available / (1024**3)
+        batch_size = 1
+        if availableGb > 14:
+            batch_size = 16
+        elif availableGb > 10:
+            batch_size = 8
+        elif availableGb > 7:
+            batch_size = 4
+    return batch_size
+
+
+@dataclass
+class TortoiseAudioConfig(Coqpit):
+    sample_rate: int = 22050
+    diffusion_sample_rate: int = 24000
+    output_sample_rate: int = 24000
+
+
+@dataclass
+class TortoiseArgs(Coqpit):
+    """A dataclass to represent Tortoise model arguments that define the model structure.
+
+    Args:
+        autoregressive_batch_size (int): The size of the auto-regressive batch.
+        enable_redaction (bool, optional): Whether to enable redaction. Defaults to True.
+        high_vram (bool, optional): Whether to use high VRAM. Defaults to False.
+        kv_cache (bool, optional): Whether to use the kv_cache. Defaults to True.
+        ar_checkpoint (str, optional): The checkpoint for the autoregressive model. Defaults to None.
+        clvp_checkpoint (str, optional): The checkpoint for the ConditionalLatentVariablePerseq model. Defaults to None.
+        diff_checkpoint (str, optional): The checkpoint for the DiffTTS model. Defaults to None.
+        num_chars (int, optional): The maximum number of characters to generate. Defaults to 255.
+        vocoder (VocType, optional): The vocoder to use for synthesis. Defaults to VocConf.Univnet.
+
+        For UnifiedVoice model:
+        ar_max_mel_tokens (int, optional): The maximum mel tokens for the autoregressive model. Defaults to 604.
+        ar_max_text_tokens (int, optional): The maximum text tokens for the autoregressive model. Defaults to 402.
+        ar_max_conditioning_inputs (int, optional): The maximum conditioning inputs for the autoregressive model. Defaults to 2.
+        ar_layers (int, optional): The number of layers for the autoregressive model. Defaults to 30.
+        ar_model_dim (int, optional): The model dimension for the autoregressive model. Defaults to 1024.
+        ar_heads (int, optional): The number of heads for the autoregressive model. Defaults to 16.
+        ar_number_text_tokens (int, optional): The number of text tokens for the autoregressive model. Defaults to 255.
+        ar_start_text_token (int, optional): The start text token for the autoregressive model. Defaults to 255.
+        ar_checkpointing (bool, optional): Whether to use checkpointing for the autoregressive model. Defaults to False.
+        ar_train_solo_embeddings (bool, optional): Whether to train embeddings for the autoregressive model. Defaults to False.
+
+        For DiffTTS model:
+        diff_model_channels (int, optional): The number of channels for the DiffTTS model. Defaults to 1024.
+        diff_num_layers (int, optional): The number of layers for the DiffTTS model. Defaults to 10.
+        diff_in_channels (int, optional): The input channels for the DiffTTS model. Defaults to 100.
+        diff_out_channels (int, optional): The output channels for the DiffTTS model. Defaults to 200.
+        diff_in_latent_channels (int, optional): The input latent channels for the DiffTTS model. Defaults to 1024.
+        diff_in_tokens (int, optional): The input tokens for the DiffTTS model. Defaults to 8193.
+        diff_dropout (int, optional): The dropout percentage for the DiffTTS model. Defaults to 0.
+        diff_use_fp16 (bool, optional): Whether to use fp16 for the DiffTTS model. Defaults to False.
+        diff_num_heads (int, optional): The number of heads for the DiffTTS model. Defaults to 16.
+        diff_layer_drop (int, optional): The layer dropout percentage for the DiffTTS model. Defaults to 0.
+        diff_unconditioned_percentage (int, optional): The percentage of unconditioned inputs for the DiffTTS model. Defaults to 0.
+
+        For ConditionalLatentVariablePerseq model:
+        clvp_dim_text (int): The dimension of the text input for the CLVP module. Defaults to 768.
+        clvp_dim_speech (int): The dimension of the speech input for the CLVP module. Defaults to 768.
+        clvp_dim_latent (int): The dimension of the latent representation for the CLVP module. Defaults to 768.
+        clvp_num_text_tokens (int): The number of text tokens used by the CLVP module. Defaults to 256.
+        clvp_text_enc_depth (int): The depth of the text encoder in the CLVP module. Defaults to 20.
+        clvp_text_seq_len (int): The maximum sequence length of the text input for the CLVP module. Defaults to 350.
+        clvp_text_heads (int): The number of attention heads used by the text encoder in the CLVP module. Defaults to 12.
+        clvp_num_speech_tokens (int): The number of speech tokens used by the CLVP module. Defaults to 8192.
+        clvp_speech_enc_depth (int): The depth of the speech encoder in the CLVP module. Defaults to 20.
+        clvp_speech_heads (int): The number of attention heads used by the speech encoder in the CLVP module. Defaults to 12.
+        clvp_speech_seq_len (int): The maximum sequence length of the speech input for the CLVP module. Defaults to 430.
+        clvp_use_xformers (bool): A flag indicating whether the model uses transformers in the CLVP module. Defaults to True.
+        duration_const (int): A constant value used in the model. Defaults to 102400.
+    """
+
+    autoregressive_batch_size: int = 1
+    enable_redaction: bool = False
+    high_vram: bool = False
+    kv_cache: bool = True
+    ar_checkpoint: str = None
+    clvp_checkpoint: str = None
+    diff_checkpoint: str = None
+    num_chars: int = 255
+    vocoder: VocType = VocConf.Univnet
+
+    # UnifiedVoice params
+    ar_max_mel_tokens: int = 604
+    ar_max_text_tokens: int = 402
+    ar_max_conditioning_inputs: int = 2
+    ar_layers: int = 30
+    ar_model_dim: int = 1024
+    ar_heads: int = 16
+    ar_number_text_tokens: int = 255
+    ar_start_text_token: int = 255
+    ar_checkpointing: bool = False
+    ar_train_solo_embeddings: bool = False
+
+    # DiffTTS params
+    diff_model_channels: int = 1024
+    diff_num_layers: int = 10
+    diff_in_channels: int = 100
+    diff_out_channels: int = 200
+    diff_in_latent_channels: int = 1024
+    diff_in_tokens: int = 8193
+    diff_dropout: int = 0
+    diff_use_fp16: bool = False
+    diff_num_heads: int = 16
+    diff_layer_drop: int = 0
+    diff_unconditioned_percentage: int = 0
+
+    # clvp params
+    clvp_dim_text: int = 768
+    clvp_dim_speech: int = 768
+    clvp_dim_latent: int = 768
+    clvp_num_text_tokens: int = 256
+    clvp_text_enc_depth: int = 20
+    clvp_text_seq_len: int = 350
+    clvp_text_heads: int = 12
+    clvp_num_speech_tokens: int = 8192
+    clvp_speech_enc_depth: int = 20
+    clvp_speech_heads: int = 12
+    clvp_speech_seq_len: int = 430
+    clvp_use_xformers: bool = True
+    # constants
+    duration_const: int = 102400
+
+
+class Tortoise(BaseTTS):
+    """Tortoise model class.
+
+    Currently only supports inference.
+
+    Examples:
+        >>> from TTS.tts.configs.tortoise_config import TortoiseConfig
+        >>> from TTS.tts.models.tortoise import Tortoise
+        >>> config = TortoiseConfig()
+        >>> model = Tortoise.inif_from_config(config)
+        >>> model.load_checkpoint(config, checkpoint_dir="paths/to/models_dir/", eval=True)
+    """
+
+    def __init__(self, config: Coqpit):
+        super().__init__(config, ap=None, tokenizer=None)
+        self.mel_norm_path = None
+        self.config = config
+        self.ar_checkpoint = self.args.ar_checkpoint
+        self.diff_checkpoint = self.args.diff_checkpoint  # TODO: check if this is even needed
+        self.models_dir = config.model_dir
+        self.autoregressive_batch_size = (
+            pick_best_batch_size_for_gpu()
+            if self.args.autoregressive_batch_size is None
+            else self.args.autoregressive_batch_size
+        )
+        self.enable_redaction = self.args.enable_redaction
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if self.enable_redaction:
+            self.aligner = Wav2VecAlignment()
+
+        self.tokenizer = VoiceBpeTokenizer()
+
+        self.autoregressive = UnifiedVoice(
+            max_mel_tokens=self.args.ar_max_mel_tokens,
+            max_text_tokens=self.args.ar_max_text_tokens,
+            max_conditioning_inputs=self.args.ar_max_conditioning_inputs,
+            layers=self.args.ar_layers,
+            model_dim=self.args.ar_model_dim,
+            heads=self.args.ar_heads,
+            number_text_tokens=self.args.ar_number_text_tokens,
+            start_text_token=self.args.ar_start_text_token,
+            checkpointing=self.args.ar_checkpointing,
+            train_solo_embeddings=self.args.ar_train_solo_embeddings,
+        ).cpu()
+
+        self.diffusion = DiffusionTts(
+            model_channels=self.args.diff_model_channels,
+            num_layers=self.args.diff_num_layers,
+            in_channels=self.args.diff_in_channels,
+            out_channels=self.args.diff_out_channels,
+            in_latent_channels=self.args.diff_in_latent_channels,
+            in_tokens=self.args.diff_in_tokens,
+            dropout=self.args.diff_dropout,
+            use_fp16=self.args.diff_use_fp16,
+            num_heads=self.args.diff_num_heads,
+            layer_drop=self.args.diff_layer_drop,
+            unconditioned_percentage=self.args.diff_unconditioned_percentage,
+        ).cpu()
+
+        self.clvp = CLVP(
+            dim_text=self.args.clvp_dim_text,
+            dim_speech=self.args.clvp_dim_speech,
+            dim_latent=self.args.clvp_dim_latent,
+            num_text_tokens=self.args.clvp_num_text_tokens,
+            text_enc_depth=self.args.clvp_text_enc_depth,
+            text_seq_len=self.args.clvp_text_seq_len,
+            text_heads=self.args.clvp_text_heads,
+            num_speech_tokens=self.args.clvp_num_speech_tokens,
+            speech_enc_depth=self.args.clvp_speech_enc_depth,
+            speech_heads=self.args.clvp_speech_heads,
+            speech_seq_len=self.args.clvp_speech_seq_len,
+            use_xformers=self.args.clvp_use_xformers,
+        ).cpu()
+
+        self.vocoder = self.args.vocoder.value.constructor().cpu()
+
+        # Random latent generators (RLGs) are loaded lazily.
+        self.rlg_auto = None
+        self.rlg_diffusion = None
+
+        if self.args.high_vram:
+            self.autoregressive = self.autoregressive.to(self.device)
+            self.diffusion = self.diffusion.to(self.device)
+            self.clvp = self.clvp.to(self.device)
+            self.vocoder = self.vocoder.to(self.device)
+        self.high_vram = self.args.high_vram
+
+    @contextmanager
+    def temporary_cuda(self, model):
+        if self.high_vram:
+            yield model
+        else:
+            m = model.to(self.device)
+            yield m
+            m = model.cpu()
+
+    def get_conditioning_latents(
+        self,
+        voice_samples,
+        return_mels=False,
+        latent_averaging_mode=0,
+        original_tortoise=False,
+    ):
+        """
+        Transforms one or more voice_samples into a tuple (autoregressive_conditioning_latent, diffusion_conditioning_latent).
+        These are expressive learned latents that encode aspects of the provided clips like voice, intonation, and acoustic
+        properties.
+        :param voice_samples: List of arbitrary reference clips, which should be *pairs* of torch tensors containing arbitrary kHz waveform data.
+        :param latent_averaging_mode: 0/1/2 for following modes:
+            0 - latents will be generated as in original tortoise, using ~4.27s from each voice sample, averaging latent across all samples
+            1 - latents will be generated using (almost) entire voice samples, averaged across all the ~4.27s chunks
+            2 - latents will be generated using (almost) entire voice samples, averaged per voice sample
+        """
+        assert latent_averaging_mode in [
+            0,
+            1,
+            2,
+        ], "latent_averaging mode has to be one of (0, 1, 2)"
+
+        with torch.no_grad():
+            voice_samples = [[v.to(self.device) for v in ls] for ls in voice_samples]
+
+            auto_conds = []
+            for ls in voice_samples:
+                auto_conds.append(format_conditioning(ls[0], device=self.device, mel_norm_file=self.mel_norm_path))
+            auto_conds = torch.stack(auto_conds, dim=1)
+            with self.temporary_cuda(self.autoregressive) as ar:
+                auto_latent = ar.get_conditioning(auto_conds)
+
+            diffusion_conds = []
+
+            DURS_CONST = self.args.duration_const
+            for ls in voice_samples:
+                # The diffuser operates at a sample rate of 24000 (except for the latent inputs)
+                sample = torchaudio.functional.resample(ls[0], 22050, 24000) if original_tortoise else ls[1]
+                if latent_averaging_mode == 0:
+                    sample = pad_or_truncate(sample, DURS_CONST)
+                    cond_mel = wav_to_univnet_mel(
+                        sample.to(self.device),
+                        do_normalization=False,
+                        device=self.device,
+                    )
+                    diffusion_conds.append(cond_mel)
+                else:
+                    from math import ceil
+
+                    if latent_averaging_mode == 2:
+                        temp_diffusion_conds = []
+                    for chunk in range(ceil(sample.shape[1] / DURS_CONST)):
+                        current_sample = sample[:, chunk * DURS_CONST : (chunk + 1) * DURS_CONST]
+                        current_sample = pad_or_truncate(current_sample, DURS_CONST)
+                        cond_mel = wav_to_univnet_mel(
+                            current_sample.to(self.device),
+                            do_normalization=False,
+                            device=self.device,
+                        )
+                        if latent_averaging_mode == 1:
+                            diffusion_conds.append(cond_mel)
+                        elif latent_averaging_mode == 2:
+                            temp_diffusion_conds.append(cond_mel)
+                    if latent_averaging_mode == 2:
+                        diffusion_conds.append(torch.stack(temp_diffusion_conds).mean(0))
+            diffusion_conds = torch.stack(diffusion_conds, dim=1)
+
+            with self.temporary_cuda(self.diffusion) as diffusion:
+                diffusion_latent = diffusion.get_conditioning(diffusion_conds)
+
+        if return_mels:
+            return auto_latent, diffusion_latent, auto_conds, diffusion_conds
+        return auto_latent, diffusion_latent
+
+    def get_random_conditioning_latents(self):
+        # Lazy-load the RLG models.
+        if self.rlg_auto is None:
+            self.rlg_auto = RandomLatentConverter(1024).eval()
+            self.rlg_auto.load_state_dict(
+                torch.load(
+                    os.path.join(self.models_dir, "rlg_auto.pth"),
+                    map_location=torch.device("cpu"),
+                )
+            )
+            self.rlg_diffusion = RandomLatentConverter(2048).eval()
+            self.rlg_diffusion.load_state_dict(
+                torch.load(
+                    os.path.join(self.models_dir, "rlg_diffuser.pth"),
+                    map_location=torch.device("cpu"),
+                )
+            )
+        with torch.no_grad():
+            return self.rlg_auto(torch.tensor([0.0])), self.rlg_diffusion(torch.tensor([0.0]))
+
+    def synthesize(self, text, config, speaker_id="random", voice_dirs=None, **kwargs):
+        """Synthesize speech with the given input text.
+
+        Args:
+            text (str): Input text.
+            config (TortoiseConfig): Config with inference parameters.
+            speaker_id (str): One of the available speaker names. If `random`, it generates a random speaker.
+            voice_dirs (List[str]): List of paths that host reference audio files for speakers. Defaults to None.
+            **kwargs: Inference settings. See `inference()`.
+
+        Returns:
+            A dictionary of the output values with `wav` as output waveform, `deterministic_seed` as seed used at inference,
+            `text_input` as text token IDs after tokenizer, `voice_samples` as samples used for cloning, `conditioning_latents`
+            as latents used at inference.
+
+        """
+
+        speaker_id = "random" if speaker_id is None else speaker_id
+
+        if voice_dirs is not None:
+            voice_dirs = [voice_dirs]
+            voice_samples, conditioning_latents = load_voice(speaker_id, voice_dirs)
+
+        else:
+            voice_samples, conditioning_latents = load_voice(speaker_id)
+
+        outputs = self.inference_with_config(
+            text, config, voice_samples=voice_samples, conditioning_latents=conditioning_latents, **kwargs
+        )
+
+        return_dict = {
+            "wav": outputs["wav"],
+            "deterministic_seed": outputs["deterministic_seed"],
+            "text_inputs": outputs["text"],
+            "voice_samples": outputs["voice_samples"],
+            "conditioning_latents": outputs["conditioning_latents"],
+        }
+
+        return return_dict
+
+    def inference_with_config(self, text, config, **kwargs):
+        """
+        inference with config
+        #TODO describe in detail
+        """
+        # Use generally found best tuning knobs for generation.
+        settings = {
+            "temperature": config.temperature,
+            "length_penalty": config.length_penalty,
+            "repetition_penalty": config.repetition_penalty,
+            "top_p": config.top_p,
+            "cond_free_k": config.cond_free_k,
+            "diffusion_temperature": config.diffusion_temperature,
+            "sampler": config.sampler,
+        }
+        # Presets are defined here.
+        presets = {
+            "single_sample": {
+                "num_autoregressive_samples": 8,
+                "diffusion_iterations": 10,
+                "sampler": "ddim",
+            },
+            "ultra_fast": {
+                "num_autoregressive_samples": 16,
+                "diffusion_iterations": 10,
+                "sampler": "ddim",
+            },
+            "ultra_fast_old": {
+                "num_autoregressive_samples": 16,
+                "diffusion_iterations": 30,
+                "cond_free": False,
+            },
+            "very_fast": {
+                "num_autoregressive_samples": 32,
+                "diffusion_iterations": 30,
+                "sampler": "dpm++2m",
+            },
+            "fast": {
+                "num_autoregressive_samples": 5,
+                "diffusion_iterations": 50,
+                "sampler": "ddim",
+            },
+            "fast_old": {"num_autoregressive_samples": 96, "diffusion_iterations": 80},
+            "standard": {
+                "num_autoregressive_samples": 5,
+                "diffusion_iterations": 200,
+            },
+            "high_quality": {
+                "num_autoregressive_samples": 256,
+                "diffusion_iterations": 400,
+            },
+        }
+        if "preset" in kwargs:
+            settings.update(presets[kwargs["preset"]])
+            kwargs.pop("preset")
+        settings.update(kwargs)  # allow overriding of preset settings with kwargs
+        return self.inference(text, **settings)
+
+    def inference(
+        self,
+        text,
+        voice_samples=None,
+        conditioning_latents=None,
+        k=1,
+        verbose=True,
+        use_deterministic_seed=None,
+        return_deterministic_state=False,
+        latent_averaging_mode=0,
+        # autoregressive generation parameters follow
+        num_autoregressive_samples=16,
+        temperature=0.8,
+        length_penalty=1,
+        repetition_penalty=2.0,
+        top_p=0.8,
+        max_mel_tokens=500,
+        # diffusion generation parameters follow
+        diffusion_iterations=100,
+        cond_free=True,
+        cond_free_k=2,
+        diffusion_temperature=1.0,
+        sampler="ddim",
+        half=True,
+        original_tortoise=False,
+        **hf_generate_kwargs,
+    ):
+        """
+        This function produces an audio clip of the given text being spoken with the given reference voice.
+
+        Args:
+            text: (str) Text to be spoken.
+            voice_samples: (List[Tuple[torch.Tensor]]) List of an arbitrary number of reference clips, which should be tuple-pairs
+                of torch tensors containing arbitrary kHz waveform data.
+            conditioning_latents: (Tuple[autoregressive_conditioning_latent, diffusion_conditioning_latent]) A tuple of
+                (autoregressive_conditioning_latent, diffusion_conditioning_latent), which can be provided in lieu
+                of voice_samples. This is ignored unless `voice_samples=None`. Conditioning latents can be retrieved
+                via `get_conditioning_latents()`.
+            k: (int) The number of returned clips. The most likely (as determined by Tortoises' CLVP model) clips are returned.
+                latent_averaging_mode: (int) 0/1/2 for following modes:
+                0 - latents will be generated as in original tortoise, using ~4.27s from each voice sample, averaging latent across all samples
+                1 - latents will be generated using (almost) entire voice samples, averaged across all the ~4.27s chunks
+                2 - latents will be generated using (almost) entire voice samples, averaged per voice sample
+            verbose: (bool) Whether or not to print log messages indicating the progress of creating a clip. Default=true.
+            num_autoregressive_samples: (int) Number of samples taken from the autoregressive model, all of which are filtered using CLVP.
+                As Tortoise is a probabilistic model, more samples means a higher probability of creating something "great".
+            temperature: (float) The softmax temperature of the autoregressive model.
+            length_penalty: (float) A length penalty applied to the autoregressive decoder. Higher settings causes the model to produce more terse outputs.
+            repetition_penalty: (float) A penalty that prevents the autoregressive decoder from repeating itself during decoding. Can be used to reduce
+                the incidence of long silences or "uhhhhhhs", etc.
+            top_p: (float) P value used in nucleus sampling. (0,1]. Lower values mean the decoder produces more "likely" (aka boring) outputs.
+            max_mel_tokens: (int) Restricts the output length. (0,600] integer. Each unit is 1/20 of a second.
+            typical_sampling: (bool) Turns typical sampling on or off. This sampling mode is discussed in this paper: https://arxiv.org/abs/2202.00666
+                I was interested in the premise, but the results were not as good as I was hoping. This is off by default, but could use some tuning.
+            typical_mass: (float) The typical_mass parameter from the typical_sampling algorithm.
+            diffusion_iterations: (int) Number of diffusion steps to perform. [0,4000]. More steps means the network has more chances to iteratively
+                refine the output, which should theoretically mean a higher quality output. Generally a value above 250 is not noticeably better, however.
+            cond_free: (bool) Whether or not to perform conditioning-free diffusion. Conditioning-free diffusion performs two forward passes for
+                each diffusion step: one with the outputs of the autoregressive model and one with no conditioning priors. The output of the two
+                is blended according to the cond_free_k value below. Conditioning-free diffusion is the real deal, and dramatically improves realism.
+            cond_free_k: (float) Knob that determines how to balance the conditioning free signal with the conditioning-present signal. [0,inf].
+                As cond_free_k increases, the output becomes dominated by the conditioning-free signal.
+            diffusion_temperature: (float) Controls the variance of the noise fed into the diffusion model. [0,1]. Values at 0
+                                      are the "mean" prediction of the diffusion network and will sound bland and smeared.
+            hf_generate_kwargs: (**kwargs) The huggingface Transformers generate API is used for the autoregressive transformer.
+                                    Extra keyword args fed to this function get forwarded directly to that API. Documentation
+                                    here: https://huggingface.co/docs/transformers/internal/generation_utils
+
+        Returns:
+            Generated audio clip(s) as a torch tensor. Shape 1,S if k=1 else, (k,1,S) where S is the sample length.
+            Sample rate is 24kHz.
+        """
+        deterministic_seed = deterministic_state(seed=use_deterministic_seed)
+
+        text_tokens = torch.IntTensor(self.tokenizer.encode(text)).unsqueeze(0).to(self.device)
+        text_tokens = F.pad(text_tokens, (0, 1))  # This may not be necessary.
+        assert (
+            text_tokens.shape[-1] < 400
+        ), "Too much text provided. Break the text up into separate segments and re-try inference."
+
+        if voice_samples is not None:
+            (
+                auto_conditioning,
+                diffusion_conditioning,
+                _,
+                _,
+            ) = self.get_conditioning_latents(
+                voice_samples,
+                return_mels=True,
+                latent_averaging_mode=latent_averaging_mode,
+                original_tortoise=original_tortoise,
+            )
+        elif conditioning_latents is not None:
+            auto_conditioning, diffusion_conditioning = conditioning_latents
+        else:
+            (
+                auto_conditioning,
+                diffusion_conditioning,
+            ) = self.get_random_conditioning_latents()
+        auto_conditioning = auto_conditioning.to(self.device)
+        diffusion_conditioning = diffusion_conditioning.to(self.device)
+
+        diffuser = load_discrete_vocoder_diffuser(
+            desired_diffusion_steps=diffusion_iterations, cond_free=cond_free, cond_free_k=cond_free_k, sampler=sampler
+        )
+
+        # in the case of single_sample,
+        orig_batch_size = self.autoregressive_batch_size
+        while num_autoregressive_samples % self.autoregressive_batch_size:
+            self.autoregressive_batch_size //= 2
+        with torch.no_grad():
+            samples = []
+            num_batches = num_autoregressive_samples // self.autoregressive_batch_size
+            stop_mel_token = self.autoregressive.stop_mel_token
+            calm_token = (
+                83  # This is the token for coding silence, which is fixed in place with "fix_autoregressive_output"
+            )
+            self.autoregressive = self.autoregressive.to(self.device)
+            if verbose:
+                print("Generating autoregressive samples..")
+            with self.temporary_cuda(self.autoregressive) as autoregressive, torch.autocast(
+                device_type="cuda", dtype=torch.float16, enabled=half
+            ):
+                for b in tqdm(range(num_batches), disable=not verbose):
+                    codes = autoregressive.inference_speech(
+                        auto_conditioning,
+                        text_tokens,
+                        do_sample=True,
+                        top_p=top_p,
+                        temperature=temperature,
+                        num_return_sequences=self.autoregressive_batch_size,
+                        length_penalty=length_penalty,
+                        repetition_penalty=repetition_penalty,
+                        max_generate_length=max_mel_tokens,
+                        **hf_generate_kwargs,
+                    )
+                    padding_needed = max_mel_tokens - codes.shape[1]
+                    codes = F.pad(codes, (0, padding_needed), value=stop_mel_token)
+                    samples.append(codes)
+            self.autoregressive_batch_size = orig_batch_size  # in the case of single_sample
+
+            clip_results = []
+            with self.temporary_cuda(self.clvp) as clvp, torch.autocast(
+                device_type="cuda", dtype=torch.float16, enabled=half
+            ):
+                for batch in tqdm(samples, disable=not verbose):
+                    for i in range(batch.shape[0]):
+                        batch[i] = fix_autoregressive_output(batch[i], stop_mel_token)
+                    clvp_res = clvp(
+                        text_tokens.repeat(batch.shape[0], 1),
+                        batch,
+                        return_loss=False,
+                    )
+                    clip_results.append(clvp_res)
+
+                clip_results = torch.cat(clip_results, dim=0)
+                samples = torch.cat(samples, dim=0)
+                best_results = samples[torch.topk(clip_results, k=k).indices]
+            del samples
+
+            # The diffusion model actually wants the last hidden layer from the autoregressive model as conditioning
+            # inputs. Re-produce those for the top results. This could be made more efficient by storing all of these
+            # results, but will increase memory usage.
+            with self.temporary_cuda(self.autoregressive) as autoregressive:
+                best_latents = autoregressive(
+                    auto_conditioning.repeat(k, 1),
+                    text_tokens.repeat(k, 1),
+                    torch.tensor([text_tokens.shape[-1]], device=text_tokens.device),
+                    best_results,
+                    torch.tensor(
+                        [best_results.shape[-1] * self.autoregressive.mel_length_compression],
+                        device=text_tokens.device,
+                    ),
+                    return_latent=True,
+                    clip_inputs=False,
+                )
+            del auto_conditioning
+
+            if verbose:
+                print("Transforming autoregressive outputs into audio..")
+            wav_candidates = []
+            for b in range(best_results.shape[0]):
+                codes = best_results[b].unsqueeze(0)
+                latents = best_latents[b].unsqueeze(0)
+
+                # Find the first occurrence of the "calm" token and trim the codes to that.
+                ctokens = 0
+                for code in range(codes.shape[-1]):
+                    if codes[0, code] == calm_token:
+                        ctokens += 1
+                    else:
+                        ctokens = 0
+                    if ctokens > 8:  # 8 tokens gives the diffusion model some "breathing room" to terminate speech.
+                        latents = latents[:, :code]
+                        break
+                with self.temporary_cuda(self.diffusion) as diffusion:
+                    mel = do_spectrogram_diffusion(
+                        diffusion,
+                        diffuser,
+                        latents,
+                        diffusion_conditioning,
+                        temperature=diffusion_temperature,
+                        verbose=verbose,
+                    )
+                with self.temporary_cuda(self.vocoder) as vocoder:
+                    wav = vocoder.inference(mel)
+                    wav_candidates.append(wav.cpu())
+
+            def potentially_redact(clip, text):
+                if self.enable_redaction:
+                    return self.aligner.redact(clip.squeeze(1), text).unsqueeze(1)
+                return clip
+
+            wav_candidates = [potentially_redact(wav_candidate, text) for wav_candidate in wav_candidates]
+
+            if len(wav_candidates) > 1:
+                res = wav_candidates
+            else:
+                res = wav_candidates[0]
+
+        return_dict = {
+            "wav": res,
+            "deterministic_seed": None,
+            "text": None,
+            "voice_samples": None,
+            "conditioning_latents": None,
+        }
+        if return_deterministic_state:
+            return_dict = {
+                "wav": res,
+                "deterministic_seed": deterministic_seed,
+                "text": text,
+                "voice_samples": voice_samples,
+                "conditioning_latents": conditioning_latents,
+            }
+        return return_dict
+
+    def forward(self):
+        raise NotImplementedError("Tortoise Training is not implemented")
+
+    def eval_step(self):
+        raise NotImplementedError("Tortoise Training is not implemented")
+
+    @staticmethod
+    def init_from_config(config: "TortoiseConfig", **kwargs):  # pylint: disable=unused-argument
+        return Tortoise(config)
+
+    def load_checkpoint(
+        self,
+        config,
+        checkpoint_dir,
+        ar_checkpoint_path=None,
+        diff_checkpoint_path=None,
+        clvp_checkpoint_path=None,
+        vocoder_checkpoint_path=None,
+        eval=False,
+        strict=True,
+        **kwargs,
+    ):  # pylint: disable=unused-argument, redefined-builtin
+        """Load a model checkpoints from a directory. This model is with multiple checkpoint files and it
+        expects to have all the files to be under the given `checkpoint_dir` with the rigth names.
+        If eval is True, set the model to eval mode.
+
+        Args:
+            config (TortoiseConfig): The model config.
+            checkpoint_dir (str): The directory where the checkpoints are stored.
+            ar_checkpoint_path (str, optional): The path to the autoregressive checkpoint. Defaults to None.
+            diff_checkpoint_path (str, optional): The path to the diffusion checkpoint. Defaults to None.
+            clvp_checkpoint_path (str, optional): The path to the CLVP checkpoint. Defaults to None.
+            vocoder_checkpoint_path (str, optional): The path to the vocoder checkpoint. Defaults to None.
+            eval (bool, optional): Whether to set the model to eval mode. Defaults to False.
+            strict (bool, optional): Whether to load the model strictly. Defaults to True.
+        """
+        if self.models_dir is None:
+            self.models_dir = checkpoint_dir
+        ar_path = ar_checkpoint_path or os.path.join(checkpoint_dir, "autoregressive.pth")
+        diff_path = diff_checkpoint_path or os.path.join(checkpoint_dir, "diffusion_decoder.pth")
+        clvp_path = clvp_checkpoint_path or os.path.join(checkpoint_dir, "clvp2.pth")
+        vocoder_checkpoint_path = vocoder_checkpoint_path or os.path.join(checkpoint_dir, "vocoder.pth")
+        self.mel_norm_path = os.path.join(checkpoint_dir, "mel_norms.pth")
+
+        if os.path.exists(ar_path):
+            # remove keys from the checkpoint that are not in the model
+            checkpoint = torch.load(ar_path, map_location=torch.device("cpu"))
+
+            # strict set False
+            # due to removed `bias` and `masked_bias` changes in Transformers
+            self.autoregressive.load_state_dict(checkpoint, strict=False)
+
+        if os.path.exists(diff_path):
+            self.diffusion.load_state_dict(torch.load(diff_path), strict=strict)
+
+        if os.path.exists(clvp_path):
+            self.clvp.load_state_dict(torch.load(clvp_path), strict=strict)
+
+        if os.path.exists(vocoder_checkpoint_path):
+            self.vocoder.load_state_dict(
+                config.model_args.vocoder.value.optionally_index(
+                    torch.load(
+                        vocoder_checkpoint_path,
+                        map_location=torch.device("cpu"),
+                    )
+                )
+            )
+
+        if eval:
+            self.autoregressive.post_init_gpt2_config(self.args.kv_cache)
+            self.autoregressive.eval()
+            self.diffusion.eval()
+            self.clvp.eval()
+            self.vocoder.eval()
+
+    def train_step(self):
+        raise NotImplementedError("Tortoise Training is not implemented")

TTS/tts/models/xtts.py

@@ -0,0 +1,791 @@
+import os
+from dataclasses import dataclass
+
+import librosa
+import torch
+import torch.nn.functional as F
+import torchaudio
+from coqpit import Coqpit
+
+from TTS.tts.layers.xtts.gpt import GPT
+from TTS.tts.layers.xtts.hifigan_decoder import HifiDecoder
+from TTS.tts.layers.xtts.stream_generator import init_stream_support
+from TTS.tts.layers.xtts.tokenizer import VoiceBpeTokenizer, split_sentence
+from TTS.tts.layers.xtts.xtts_manager import SpeakerManager, LanguageManager
+from TTS.tts.models.base_tts import BaseTTS
+from TTS.utils.io import load_fsspec
+
+init_stream_support()
+
+
+def wav_to_mel_cloning(
+    wav,
+    mel_norms_file="../experiments/clips_mel_norms.pth",
+    mel_norms=None,
+    device=torch.device("cpu"),
+    n_fft=4096,
+    hop_length=1024,
+    win_length=4096,
+    power=2,
+    normalized=False,
+    sample_rate=22050,
+    f_min=0,
+    f_max=8000,
+    n_mels=80,
+):
+    """
+    Convert waveform to mel-spectrogram with hard-coded parameters for cloning.
+
+    Args:
+        wav (torch.Tensor): Input waveform tensor.
+        mel_norms_file (str): Path to mel-spectrogram normalization file.
+        mel_norms (torch.Tensor): Mel-spectrogram normalization tensor.
+        device (torch.device): Device to use for computation.
+
+    Returns:
+        torch.Tensor: Mel-spectrogram tensor.
+    """
+    mel_stft = torchaudio.transforms.MelSpectrogram(
+        n_fft=n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        power=power,
+        normalized=normalized,
+        sample_rate=sample_rate,
+        f_min=f_min,
+        f_max=f_max,
+        n_mels=n_mels,
+        norm="slaney",
+    ).to(device)
+    wav = wav.to(device)
+    mel = mel_stft(wav)
+    mel = torch.log(torch.clamp(mel, min=1e-5))
+    if mel_norms is None:
+        mel_norms = torch.load(mel_norms_file, map_location=device)
+    mel = mel / mel_norms.unsqueeze(0).unsqueeze(-1)
+    return mel
+
+
+def load_audio(audiopath, sampling_rate):
+    # better load setting following: https://github.com/faroit/python_audio_loading_benchmark
+
+    # torchaudio should chose proper backend to load audio depending on platform
+    audio, lsr = torchaudio.load(audiopath)
+
+    # stereo to mono if needed
+    if audio.size(0) != 1:
+        audio = torch.mean(audio, dim=0, keepdim=True)
+
+    if lsr != sampling_rate:
+        audio = torchaudio.functional.resample(audio, lsr, sampling_rate)
+
+    # Check some assumptions about audio range. This should be automatically fixed in load_wav_to_torch, but might not be in some edge cases, where we should squawk.
+    # '10' is arbitrarily chosen since it seems like audio will often "overdrive" the [-1,1] bounds.
+    if torch.any(audio > 10) or not torch.any(audio < 0):
+        print(f"Error with {audiopath}. Max={audio.max()} min={audio.min()}")
+    # clip audio invalid values
+    audio.clip_(-1, 1)
+    return audio
+
+
+def pad_or_truncate(t, length):
+    """
+    Ensure a given tensor t has a specified sequence length by either padding it with zeros or clipping it.
+
+    Args:
+        t (torch.Tensor): The input tensor to be padded or truncated.
+        length (int): The desired length of the tensor.
+
+    Returns:
+        torch.Tensor: The padded or truncated tensor.
+    """
+    tp = t[..., :length]
+    if t.shape[-1] == length:
+        tp = t
+    elif t.shape[-1] < length:
+        tp = F.pad(t, (0, length - t.shape[-1]))
+    return tp
+
+
+@dataclass
+class XttsAudioConfig(Coqpit):
+    """
+    Configuration class for audio-related parameters in the XTTS model.
+
+    Args:
+        sample_rate (int): The sample rate in which the GPT operates.
+        output_sample_rate (int): The sample rate of the output audio waveform.
+    """
+
+    sample_rate: int = 22050
+    output_sample_rate: int = 24000
+
+
+@dataclass
+class XttsArgs(Coqpit):
+    """A dataclass to represent XTTS model arguments that define the model structure.
+
+    Args:
+        gpt_batch_size (int): The size of the auto-regressive batch.
+        enable_redaction (bool, optional): Whether to enable redaction. Defaults to True.
+        kv_cache (bool, optional): Whether to use the kv_cache. Defaults to True.
+        gpt_checkpoint (str, optional): The checkpoint for the autoregressive model. Defaults to None.
+        clvp_checkpoint (str, optional): The checkpoint for the ConditionalLatentVariablePerseq model. Defaults to None.
+        decoder_checkpoint (str, optional): The checkpoint for the DiffTTS model. Defaults to None.
+        num_chars (int, optional): The maximum number of characters to generate. Defaults to 255.
+
+        For GPT model:
+        gpt_max_audio_tokens (int, optional): The maximum mel tokens for the autoregressive model. Defaults to 604.
+        gpt_max_text_tokens (int, optional): The maximum text tokens for the autoregressive model. Defaults to 402.
+        gpt_max_prompt_tokens (int, optional): The maximum prompt tokens or the autoregressive model. Defaults to 70.
+        gpt_layers (int, optional): The number of layers for the autoregressive model. Defaults to 30.
+        gpt_n_model_channels (int, optional): The model dimension for the autoregressive model. Defaults to 1024.
+        gpt_n_heads (int, optional): The number of heads for the autoregressive model. Defaults to 16.
+        gpt_number_text_tokens (int, optional): The number of text tokens for the autoregressive model. Defaults to 255.
+        gpt_start_text_token (int, optional): The start text token for the autoregressive model. Defaults to 255.
+        gpt_checkpointing (bool, optional): Whether to use checkpointing for the autoregressive model. Defaults to False.
+        gpt_train_solo_embeddings (bool, optional): Whether to train embeddings for the autoregressive model. Defaults to False.
+        gpt_code_stride_len (int, optional): The hop_size of dvae and consequently of the gpt output. Defaults to 1024.
+        gpt_use_masking_gt_prompt_approach (bool, optional):  If True, it will use ground truth as prompt and it will mask the loss to avoid repetition. Defaults to True.
+        gpt_use_perceiver_resampler (bool, optional):  If True, it will use perceiver resampler from flamingo paper - https://arxiv.org/abs/2204.14198. Defaults to False.
+    """
+
+    gpt_batch_size: int = 1
+    enable_redaction: bool = False
+    kv_cache: bool = True
+    gpt_checkpoint: str = None
+    clvp_checkpoint: str = None
+    decoder_checkpoint: str = None
+    num_chars: int = 255
+
+    # XTTS GPT Encoder params
+    tokenizer_file: str = ""
+    gpt_max_audio_tokens: int = 605
+    gpt_max_text_tokens: int = 402
+    gpt_max_prompt_tokens: int = 70
+    gpt_layers: int = 30
+    gpt_n_model_channels: int = 1024
+    gpt_n_heads: int = 16
+    gpt_number_text_tokens: int = None
+    gpt_start_text_token: int = None
+    gpt_stop_text_token: int = None
+    gpt_num_audio_tokens: int = 8194
+    gpt_start_audio_token: int = 8192
+    gpt_stop_audio_token: int = 8193
+    gpt_code_stride_len: int = 1024
+    gpt_use_masking_gt_prompt_approach: bool = True
+    gpt_use_perceiver_resampler: bool = False
+
+    # HifiGAN Decoder params
+    input_sample_rate: int = 22050
+    output_sample_rate: int = 24000
+    output_hop_length: int = 256
+    decoder_input_dim: int = 1024
+    d_vector_dim: int = 512
+    cond_d_vector_in_each_upsampling_layer: bool = True
+
+    # constants
+    duration_const: int = 102400
+
+
+class Xtts(BaseTTS):
+    """ⓍTTS model implementation.
+
+    ❗ Currently it only supports inference.
+
+    Examples:
+        >>> from TTS.tts.configs.xtts_config import XttsConfig
+        >>> from TTS.tts.models.xtts import Xtts
+        >>> config = XttsConfig()
+        >>> model = Xtts.inif_from_config(config)
+        >>> model.load_checkpoint(config, checkpoint_dir="paths/to/models_dir/", eval=True)
+    """
+
+    def __init__(self, config: Coqpit):
+        super().__init__(config, ap=None, tokenizer=None)
+        self.mel_stats_path = None
+        self.config = config
+        self.gpt_checkpoint = self.args.gpt_checkpoint
+        self.decoder_checkpoint = self.args.decoder_checkpoint  # TODO: check if this is even needed
+        self.models_dir = config.model_dir
+        self.gpt_batch_size = self.args.gpt_batch_size
+
+        self.tokenizer = VoiceBpeTokenizer()
+        self.gpt = None
+        self.init_models()
+        self.register_buffer("mel_stats", torch.ones(80))
+
+    def init_models(self):
+        """Initialize the models. We do it here since we need to load the tokenizer first."""
+        if self.tokenizer.tokenizer is not None:
+            self.args.gpt_number_text_tokens = self.tokenizer.get_number_tokens()
+            self.args.gpt_start_text_token = self.tokenizer.tokenizer.token_to_id("[START]")
+            self.args.gpt_stop_text_token = self.tokenizer.tokenizer.token_to_id("[STOP]")
+
+        if self.args.gpt_number_text_tokens:
+            self.gpt = GPT(
+                layers=self.args.gpt_layers,
+                model_dim=self.args.gpt_n_model_channels,
+                start_text_token=self.args.gpt_start_text_token,
+                stop_text_token=self.args.gpt_stop_text_token,
+                heads=self.args.gpt_n_heads,
+                max_text_tokens=self.args.gpt_max_text_tokens,
+                max_mel_tokens=self.args.gpt_max_audio_tokens,
+                max_prompt_tokens=self.args.gpt_max_prompt_tokens,
+                number_text_tokens=self.args.gpt_number_text_tokens,
+                num_audio_tokens=self.args.gpt_num_audio_tokens,
+                start_audio_token=self.args.gpt_start_audio_token,
+                stop_audio_token=self.args.gpt_stop_audio_token,
+                use_perceiver_resampler=self.args.gpt_use_perceiver_resampler,
+                code_stride_len=self.args.gpt_code_stride_len,
+            )
+
+        self.hifigan_decoder = HifiDecoder(
+            input_sample_rate=self.args.input_sample_rate,
+            output_sample_rate=self.args.output_sample_rate,
+            output_hop_length=self.args.output_hop_length,
+            ar_mel_length_compression=self.args.gpt_code_stride_len,
+            decoder_input_dim=self.args.decoder_input_dim,
+            d_vector_dim=self.args.d_vector_dim,
+            cond_d_vector_in_each_upsampling_layer=self.args.cond_d_vector_in_each_upsampling_layer,
+        )
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @torch.inference_mode()
+    def get_gpt_cond_latents(self, audio, sr, length: int = 30, chunk_length: int = 6):
+        """Compute the conditioning latents for the GPT model from the given audio.
+
+        Args:
+            audio (tensor): audio tensor.
+            sr (int): Sample rate of the audio.
+            length (int): Length of the audio in seconds. If < 0, use the whole audio. Defaults to 30.
+            chunk_length (int): Length of the audio chunks in seconds. When `length == chunk_length`, the whole audio
+                is being used without chunking. It must be < `length`. Defaults to 6.
+        """
+        if sr != 22050:
+            audio = torchaudio.functional.resample(audio, sr, 22050)
+        if length > 0:
+            audio = audio[:, : 22050 * length]
+        if self.args.gpt_use_perceiver_resampler:
+            style_embs = []
+            for i in range(0, audio.shape[1], 22050 * chunk_length):
+                audio_chunk = audio[:, i : i + 22050 * chunk_length]
+
+                # if the chunk is too short ignore it 
+                if audio_chunk.size(-1) < 22050 * 0.33:
+                    continue
+
+                mel_chunk = wav_to_mel_cloning(
+                    audio_chunk,
+                    mel_norms=self.mel_stats.cpu(),
+                    n_fft=2048,
+                    hop_length=256,
+                    win_length=1024,
+                    power=2,
+                    normalized=False,
+                    sample_rate=22050,
+                    f_min=0,
+                    f_max=8000,
+                    n_mels=80,
+                )
+                style_emb = self.gpt.get_style_emb(mel_chunk.to(self.device), None)
+                style_embs.append(style_emb)
+
+            # mean style embedding
+            cond_latent = torch.stack(style_embs).mean(dim=0)
+        else:
+            mel = wav_to_mel_cloning(
+                audio,
+                mel_norms=self.mel_stats.cpu(),
+                n_fft=4096,
+                hop_length=1024,
+                win_length=4096,
+                power=2,
+                normalized=False,
+                sample_rate=22050,
+                f_min=0,
+                f_max=8000,
+                n_mels=80,
+            )
+            cond_latent = self.gpt.get_style_emb(mel.to(self.device))
+        return cond_latent.transpose(1, 2)
+
+    @torch.inference_mode()
+    def get_speaker_embedding(self, audio, sr):
+        audio_16k = torchaudio.functional.resample(audio, sr, 16000)
+        return (
+            self.hifigan_decoder.speaker_encoder.forward(audio_16k.to(self.device), l2_norm=True)
+            .unsqueeze(-1)
+            .to(self.device)
+        )
+
+    @torch.inference_mode()
+    def get_conditioning_latents(
+        self,
+        audio_path,
+        max_ref_length=30,
+        gpt_cond_len=6,
+        gpt_cond_chunk_len=6,
+        librosa_trim_db=None,
+        sound_norm_refs=False,
+        load_sr=22050,
+    ):
+        """Get the conditioning latents for the GPT model from the given audio.
+
+        Args:
+            audio_path (str or List[str]): Path to reference audio file(s).
+            max_ref_length (int): Maximum length of each reference audio in seconds. Defaults to 30.
+            gpt_cond_len (int): Length of the audio used for gpt latents. Defaults to 6.
+            gpt_cond_chunk_len (int): Chunk length used for gpt latents. It must be <= gpt_conf_len. Defaults to 6.
+            librosa_trim_db (int, optional): Trim the audio using this value. If None, not trimming. Defaults to None.
+            sound_norm_refs (bool, optional): Whether to normalize the audio. Defaults to False.
+            load_sr (int, optional): Sample rate to load the audio. Defaults to 24000.
+        """
+        # deal with multiples references
+        if not isinstance(audio_path, list):
+            audio_paths = [audio_path]
+        else:
+            audio_paths = audio_path
+
+        speaker_embeddings = []
+        audios = []
+        speaker_embedding = None
+        for file_path in audio_paths:
+            audio = load_audio(file_path, load_sr)
+            audio = audio[:, : load_sr * max_ref_length].to(self.device)
+            if sound_norm_refs:
+                audio = (audio / torch.abs(audio).max()) * 0.75
+            if librosa_trim_db is not None:
+                audio = librosa.effects.trim(audio, top_db=librosa_trim_db)[0]
+
+            # compute latents for the decoder
+            speaker_embedding = self.get_speaker_embedding(audio, load_sr)
+            speaker_embeddings.append(speaker_embedding)
+
+            audios.append(audio)
+
+        # merge all the audios and compute the latents for the gpt
+        full_audio = torch.cat(audios, dim=-1)
+        gpt_cond_latents = self.get_gpt_cond_latents(
+            full_audio, load_sr, length=gpt_cond_len, chunk_length=gpt_cond_chunk_len
+        )  # [1, 1024, T]
+
+        if speaker_embeddings:
+            speaker_embedding = torch.stack(speaker_embeddings)
+            speaker_embedding = speaker_embedding.mean(dim=0)
+
+        return gpt_cond_latents, speaker_embedding
+
+    def synthesize(self, text, config, speaker_wav, language, speaker_id=None, **kwargs):
+        """Synthesize speech with the given input text.
+
+        Args:
+            text (str): Input text.
+            config (XttsConfig): Config with inference parameters.
+            speaker_wav (list): List of paths to the speaker audio files to be used for cloning.
+            language (str): Language ID of the speaker.
+            **kwargs: Inference settings. See `inference()`.
+
+        Returns:
+            A dictionary of the output values with `wav` as output waveform, `deterministic_seed` as seed used at inference,
+            `text_input` as text token IDs after tokenizer, `voice_samples` as samples used for cloning, `conditioning_latents`
+            as latents used at inference.
+
+        """
+        assert (
+            "zh-cn" if language == "zh" else language in self.config.languages
+        ), f" ❗ Language {language} is not supported. Supported languages are {self.config.languages}"
+        # Use generally found best tuning knobs for generation.
+        settings = {
+            "temperature": config.temperature,
+            "length_penalty": config.length_penalty,
+            "repetition_penalty": config.repetition_penalty,
+            "top_k": config.top_k,
+            "top_p": config.top_p,
+        }
+        settings.update(kwargs)  # allow overriding of preset settings with kwargs
+        if speaker_id is not None:
+            gpt_cond_latent, speaker_embedding = self.speaker_manager.speakers[speaker_id].values()
+            return self.inference(text, language, gpt_cond_latent, speaker_embedding, **settings)
+        settings.update({
+            "gpt_cond_len": config.gpt_cond_len,
+            "gpt_cond_chunk_len": config.gpt_cond_chunk_len,
+            "max_ref_len": config.max_ref_len,
+            "sound_norm_refs": config.sound_norm_refs,
+        })
+        return self.full_inference(text, speaker_wav, language, **settings)
+
+    @torch.inference_mode()
+    def full_inference(
+        self,
+        text,
+        ref_audio_path,
+        language,
+        # GPT inference
+        temperature=0.75,
+        length_penalty=1.0,
+        repetition_penalty=10.0,
+        top_k=50,
+        top_p=0.85,
+        do_sample=True,
+        # Cloning
+        gpt_cond_len=30,
+        gpt_cond_chunk_len=6,
+        max_ref_len=10,
+        sound_norm_refs=False,
+        **hf_generate_kwargs,
+    ):
+        """
+        This function produces an audio clip of the given text being spoken with the given reference voice.
+
+        Args:
+            text: (str) Text to be spoken.
+
+            ref_audio_path: (str) Path to a reference audio file to be used for cloning. This audio file should be >3
+                seconds long.
+
+            language: (str) Language of the voice to be generated.
+
+            temperature: (float) The softmax temperature of the autoregressive model. Defaults to 0.65.
+
+            length_penalty: (float) A length penalty applied to the autoregressive decoder. Higher settings causes the
+                model to produce more terse outputs. Defaults to 1.0.
+
+            repetition_penalty: (float) A penalty that prevents the autoregressive decoder from repeating itself during
+                decoding. Can be used to reduce the incidence of long silences or "uhhhhhhs", etc. Defaults to 2.0.
+
+            top_k: (int) K value used in top-k sampling. [0,inf]. Lower values mean the decoder produces more "likely"
+                (aka boring) outputs. Defaults to 50.
+
+            top_p: (float) P value used in nucleus sampling. (0,1]. Lower values mean the decoder produces more "likely"
+                (aka boring) outputs. Defaults to 0.8.
+
+            gpt_cond_len: (int) Length of the audio used for cloning. If audio is shorter, then audio length is used
+                else the first `gpt_cond_len` secs is used. Defaults to 30 seconds.
+
+            gpt_cond_chunk_len: (int) Chunk length used for cloning. It must be <= `gpt_cond_len`.
+                If gpt_cond_len == gpt_cond_chunk_len, no chunking. Defaults to 6 seconds.
+
+            hf_generate_kwargs: (**kwargs) The huggingface Transformers generate API is used for the autoregressive
+                transformer. Extra keyword args fed to this function get forwarded directly to that API. Documentation
+                here: https://huggingface.co/docs/transformers/internal/generation_utils
+
+        Returns:
+            Generated audio clip(s) as a torch tensor. Shape 1,S if k=1 else, (k,1,S) where S is the sample length.
+            Sample rate is 24kHz.
+        """
+        (gpt_cond_latent, speaker_embedding) = self.get_conditioning_latents(
+            audio_path=ref_audio_path,
+            gpt_cond_len=gpt_cond_len,
+            gpt_cond_chunk_len=gpt_cond_chunk_len,
+            max_ref_length=max_ref_len,
+            sound_norm_refs=sound_norm_refs,
+        )
+
+        return self.inference(
+            text,
+            language,
+            gpt_cond_latent,
+            speaker_embedding,
+            temperature=temperature,
+            length_penalty=length_penalty,
+            repetition_penalty=repetition_penalty,
+            top_k=top_k,
+            top_p=top_p,
+            do_sample=do_sample,
+            **hf_generate_kwargs,
+        )
+
+    @torch.inference_mode()
+    def inference(
+        self,
+        text,
+        language,
+        gpt_cond_latent,
+        speaker_embedding,
+        # GPT inference
+        temperature=0.75,
+        length_penalty=1.0,
+        repetition_penalty=10.0,
+        top_k=50,
+        top_p=0.85,
+        do_sample=True,
+        num_beams=1,
+        speed=1.0,
+        enable_text_splitting=False,
+        **hf_generate_kwargs,
+    ):
+        language = language.split("-")[0]  # remove the country code
+        length_scale = 1.0 / max(speed, 0.05)
+        gpt_cond_latent = gpt_cond_latent.to(self.device)
+        speaker_embedding = speaker_embedding.to(self.device)
+        if enable_text_splitting:
+            text = split_sentence(text, language, self.tokenizer.char_limits[language])
+        else:
+            text = [text]
+
+        wavs = []
+        gpt_latents_list = []
+        for sent in text:
+            sent = sent.strip().lower()
+            text_tokens = torch.IntTensor(self.tokenizer.encode(sent, lang=language)).unsqueeze(0).to(self.device)
+
+            assert (
+                text_tokens.shape[-1] < self.args.gpt_max_text_tokens
+            ), " ❗ XTTS can only generate text with a maximum of 400 tokens."
+
+            with torch.no_grad():
+                gpt_codes = self.gpt.generate(
+                    cond_latents=gpt_cond_latent,
+                    text_inputs=text_tokens,
+                    input_tokens=None,
+                    do_sample=do_sample,
+                    top_p=top_p,
+                    top_k=top_k,
+                    temperature=temperature,
+                    num_return_sequences=self.gpt_batch_size,
+                    num_beams=num_beams,
+                    length_penalty=length_penalty,
+                    repetition_penalty=repetition_penalty,
+                    output_attentions=False,
+                    **hf_generate_kwargs,
+                )
+                expected_output_len = torch.tensor(
+                    [gpt_codes.shape[-1] * self.gpt.code_stride_len], device=text_tokens.device
+                )
+
+                text_len = torch.tensor([text_tokens.shape[-1]], device=self.device)
+                gpt_latents = self.gpt(
+                    text_tokens,
+                    text_len,
+                    gpt_codes,
+                    expected_output_len,
+                    cond_latents=gpt_cond_latent,
+                    return_attentions=False,
+                    return_latent=True,
+                )
+
+                if length_scale != 1.0:
+                    gpt_latents = F.interpolate(
+                        gpt_latents.transpose(1, 2), scale_factor=length_scale, mode="linear"
+                    ).transpose(1, 2)
+
+                gpt_latents_list.append(gpt_latents.cpu())
+                wavs.append(self.hifigan_decoder(gpt_latents, g=speaker_embedding).cpu().squeeze())
+
+        return {
+            "wav": torch.cat(wavs, dim=0).numpy(),
+            "gpt_latents": torch.cat(gpt_latents_list, dim=1).numpy(),
+            "speaker_embedding": speaker_embedding,
+        }
+
+    def handle_chunks(self, wav_gen, wav_gen_prev, wav_overlap, overlap_len):
+        """Handle chunk formatting in streaming mode"""
+        wav_chunk = wav_gen[:-overlap_len]
+        if wav_gen_prev is not None:
+            wav_chunk = wav_gen[(wav_gen_prev.shape[0] - overlap_len) : -overlap_len]
+        if wav_overlap is not None:
+            # cross fade the overlap section
+            if overlap_len > len(wav_chunk):
+                # wav_chunk is smaller than overlap_len, pass on last wav_gen
+                if wav_gen_prev is not None:
+                    wav_chunk = wav_gen[(wav_gen_prev.shape[0] - overlap_len) :]
+                else:
+                    # not expecting will hit here as problem happens on last chunk
+                    wav_chunk = wav_gen[-overlap_len:]
+                return wav_chunk, wav_gen, None
+            else:
+                crossfade_wav = wav_chunk[:overlap_len]
+                crossfade_wav = crossfade_wav * torch.linspace(0.0, 1.0, overlap_len).to(crossfade_wav.device)
+                wav_chunk[:overlap_len] = wav_overlap * torch.linspace(1.0, 0.0, overlap_len).to(wav_overlap.device)
+                wav_chunk[:overlap_len] += crossfade_wav
+
+        wav_overlap = wav_gen[-overlap_len:]
+        wav_gen_prev = wav_gen
+        return wav_chunk, wav_gen_prev, wav_overlap
+
+    @torch.inference_mode()
+    def inference_stream(
+        self,
+        text,
+        language,
+        gpt_cond_latent,
+        speaker_embedding,
+        # Streaming
+        stream_chunk_size=20,
+        overlap_wav_len=1024,
+        # GPT inference
+        temperature=0.75,
+        length_penalty=1.0,
+        repetition_penalty=10.0,
+        top_k=50,
+        top_p=0.85,
+        do_sample=True,
+        speed=1.0,
+        enable_text_splitting=False,
+        **hf_generate_kwargs,
+    ):
+        language = language.split("-")[0]  # remove the country code
+        length_scale = 1.0 / max(speed, 0.05)
+        gpt_cond_latent = gpt_cond_latent.to(self.device)
+        speaker_embedding = speaker_embedding.to(self.device)
+        if enable_text_splitting:
+            text = split_sentence(text, language, self.tokenizer.char_limits[language])
+        else:
+            text = [text]
+
+        for sent in text:
+            sent = sent.strip().lower()
+            text_tokens = torch.IntTensor(self.tokenizer.encode(sent, lang=language)).unsqueeze(0).to(self.device)
+
+            assert (
+                text_tokens.shape[-1] < self.args.gpt_max_text_tokens
+            ), " ❗ XTTS can only generate text with a maximum of 400 tokens."
+
+            fake_inputs = self.gpt.compute_embeddings(
+                gpt_cond_latent.to(self.device),
+                text_tokens,
+            )
+            gpt_generator = self.gpt.get_generator(
+                fake_inputs=fake_inputs,
+                top_k=top_k,
+                top_p=top_p,
+                temperature=temperature,
+                do_sample=do_sample,
+                num_beams=1,
+                num_return_sequences=1,
+                length_penalty=float(length_penalty),
+                repetition_penalty=float(repetition_penalty),
+                output_attentions=False,
+                output_hidden_states=True,
+                **hf_generate_kwargs,
+            )
+
+            last_tokens = []
+            all_latents = []
+            wav_gen_prev = None
+            wav_overlap = None
+            is_end = False
+
+            while not is_end:
+                try:
+                    x, latent = next(gpt_generator)
+                    last_tokens += [x]
+                    all_latents += [latent]
+                except StopIteration:
+                    is_end = True
+
+                if is_end or (stream_chunk_size > 0 and len(last_tokens) >= stream_chunk_size):
+                    gpt_latents = torch.cat(all_latents, dim=0)[None, :]
+                    if length_scale != 1.0:
+                        gpt_latents = F.interpolate(
+                            gpt_latents.transpose(1, 2), scale_factor=length_scale, mode="linear"
+                        ).transpose(1, 2)
+                    wav_gen = self.hifigan_decoder(gpt_latents, g=speaker_embedding.to(self.device))
+                    wav_chunk, wav_gen_prev, wav_overlap = self.handle_chunks(
+                        wav_gen.squeeze(), wav_gen_prev, wav_overlap, overlap_wav_len
+                    )
+                    last_tokens = []
+                    yield wav_chunk
+
+    def forward(self):
+        raise NotImplementedError(
+            "XTTS has a dedicated trainer, please check the XTTS docs: https://tts.readthedocs.io/en/dev/models/xtts.html#training"
+        )
+
+    def eval_step(self):
+        raise NotImplementedError(
+            "XTTS has a dedicated trainer, please check the XTTS docs: https://tts.readthedocs.io/en/dev/models/xtts.html#training"
+        )
+
+    @staticmethod
+    def init_from_config(config: "XttsConfig", **kwargs):  # pylint: disable=unused-argument
+        return Xtts(config)
+
+    def eval(self):  # pylint: disable=redefined-builtin
+        """Sets the model to evaluation mode. Overrides the default eval() method to also set the GPT model to eval mode."""
+        self.gpt.init_gpt_for_inference()
+        super().eval()
+
+    def get_compatible_checkpoint_state_dict(self, model_path):
+        checkpoint = load_fsspec(model_path, map_location=torch.device("cpu"))["model"]
+        # remove xtts gpt trainer extra keys
+        ignore_keys = ["torch_mel_spectrogram_style_encoder", "torch_mel_spectrogram_dvae", "dvae"]
+        for key in list(checkpoint.keys()):
+            # check if it is from the coqui Trainer if so convert it
+            if key.startswith("xtts."):
+                new_key = key.replace("xtts.", "")
+                checkpoint[new_key] = checkpoint[key]
+                del checkpoint[key]
+                key = new_key
+
+            # remove unused keys
+            if key.split(".")[0] in ignore_keys:
+                del checkpoint[key]
+
+        return checkpoint
+
+    def load_checkpoint(
+        self,
+        config,
+        checkpoint_dir=None,
+        checkpoint_path=None,
+        vocab_path=None,
+        eval=True,
+        strict=True,
+        use_deepspeed=False,
+        speaker_file_path=None,
+    ):
+        """
+        Loads a checkpoint from disk and initializes the model's state and tokenizer.
+
+        Args:
+            config (dict): The configuration dictionary for the model.
+            checkpoint_dir (str, optional): The directory where the checkpoint is stored. Defaults to None.
+            checkpoint_path (str, optional): The path to the checkpoint file. Defaults to None.
+            vocab_path (str, optional): The path to the vocabulary file. Defaults to None.
+            eval (bool, optional): Whether to set the model to evaluation mode. Defaults to True.
+            strict (bool, optional): Whether to strictly enforce that the keys in the checkpoint match the keys in the model. Defaults to True.
+
+        Returns:
+            None
+        """
+
+        model_path = checkpoint_path or os.path.join(checkpoint_dir, "model.pth")
+        vocab_path = vocab_path or os.path.join(checkpoint_dir, "vocab.json")
+
+        if speaker_file_path is None and checkpoint_dir is not None:
+            speaker_file_path = os.path.join(checkpoint_dir, "speakers_xtts.pth")
+
+        self.language_manager = LanguageManager(config)
+        self.speaker_manager = None
+        if speaker_file_path is not None and os.path.exists(speaker_file_path):
+            self.speaker_manager = SpeakerManager(speaker_file_path)
+
+        if os.path.exists(vocab_path):
+            self.tokenizer = VoiceBpeTokenizer(vocab_file=vocab_path)
+
+        self.init_models()
+
+        checkpoint = self.get_compatible_checkpoint_state_dict(model_path)
+
+        # deal with v1 and v1.1. V1 has the init_gpt_for_inference keys, v1.1 do not
+        try:
+            self.load_state_dict(checkpoint, strict=strict)
+        except:
+            if eval:
+                self.gpt.init_gpt_for_inference(kv_cache=self.args.kv_cache)
+            self.load_state_dict(checkpoint, strict=strict)
+
+        if eval:
+            self.hifigan_decoder.eval()
+            self.gpt.init_gpt_for_inference(kv_cache=self.args.kv_cache, use_deepspeed=use_deepspeed)
+            self.gpt.eval()
+
+    def train_step(self):
+        raise NotImplementedError(
+            "XTTS has a dedicated trainer, please check the XTTS docs: https://tts.readthedocs.io/en/dev/models/xtts.html#training"
+        )