Add files via upload

TTS/tts/layers/__init__.py

@@ -0,0 +1 @@
+from TTS.tts.layers.losses import *

TTS/tts/layers/align_tts/duration_predictor.py

@@ -0,0 +1,21 @@
+from torch import nn
+
+from TTS.tts.layers.generic.pos_encoding import PositionalEncoding
+from TTS.tts.layers.generic.transformer import FFTransformerBlock
+
+
+class DurationPredictor(nn.Module):
+    def __init__(self, num_chars, hidden_channels, hidden_channels_ffn, num_heads):
+        super().__init__()
+        self.embed = nn.Embedding(num_chars, hidden_channels)
+        self.pos_enc = PositionalEncoding(hidden_channels, dropout_p=0.1)
+        self.FFT = FFTransformerBlock(hidden_channels, num_heads, hidden_channels_ffn, 2, 0.1)
+        self.out_layer = nn.Conv1d(hidden_channels, 1, 1)
+
+    def forward(self, text, text_lengths):
+        # B, L -> B, L
+        emb = self.embed(text)
+        emb = self.pos_enc(emb.transpose(1, 2))
+        x = self.FFT(emb, text_lengths)
+        x = self.out_layer(x).squeeze(-1)
+        return x

TTS/tts/layers/align_tts/mdn.py

@@ -0,0 +1,30 @@
+from torch import nn
+
+
+class MDNBlock(nn.Module):
+    """Mixture of Density Network implementation
+    https://arxiv.org/pdf/2003.01950.pdf
+    """
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.out_channels = out_channels
+        self.conv1 = nn.Conv1d(in_channels, in_channels, 1)
+        self.norm = nn.LayerNorm(in_channels)
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(0.1)
+        self.conv2 = nn.Conv1d(in_channels, out_channels, 1)
+
+    def forward(self, x):
+        o = self.conv1(x)
+        o = o.transpose(1, 2)
+        o = self.norm(o)
+        o = o.transpose(1, 2)
+        o = self.relu(o)
+        o = self.dropout(o)
+        mu_sigma = self.conv2(o)
+        # TODO: check this sigmoid
+        # mu = torch.sigmoid(mu_sigma[:, :self.out_channels//2, :])
+        mu = mu_sigma[:, : self.out_channels // 2, :]
+        log_sigma = mu_sigma[:, self.out_channels // 2 :, :]
+        return mu, log_sigma

TTS/tts/layers/bark/hubert/hubert_manager.py

@@ -0,0 +1,35 @@
+# From https://github.com/gitmylo/bark-voice-cloning-HuBERT-quantizer
+
+import os.path
+import shutil
+import urllib.request
+
+import huggingface_hub
+
+
+class HubertManager:
+    @staticmethod
+    def make_sure_hubert_installed(
+        download_url: str = "https://dl.fbaipublicfiles.com/hubert/hubert_base_ls960.pt", model_path: str = ""
+    ):
+        if not os.path.isfile(model_path):
+            print("Downloading HuBERT base model")
+            urllib.request.urlretrieve(download_url, model_path)
+            print("Downloaded HuBERT")
+            return model_path
+        return None
+
+    @staticmethod
+    def make_sure_tokenizer_installed(
+        model: str = "quantifier_hubert_base_ls960_14.pth",
+        repo: str = "GitMylo/bark-voice-cloning",
+        model_path: str = "",
+    ):
+        model_dir = os.path.dirname(model_path)
+        if not os.path.isfile(model_path):
+            print("Downloading HuBERT custom tokenizer")
+            huggingface_hub.hf_hub_download(repo, model, local_dir=model_dir, local_dir_use_symlinks=False)
+            shutil.move(os.path.join(model_dir, model), model_path)
+            print("Downloaded tokenizer")
+            return model_path
+        return None

TTS/tts/layers/bark/hubert/kmeans_hubert.py

@@ -0,0 +1,82 @@
+"""
+Modified HuBERT model without kmeans.
+Original author: https://github.com/lucidrains/
+Modified by: https://www.github.com/gitmylo/
+License: MIT
+"""
+
+# Modified code from https://github.com/lucidrains/audiolm-pytorch/blob/main/audiolm_pytorch/hubert_kmeans.py
+
+import logging
+from pathlib import Path
+
+import torch
+from einops import pack, unpack
+from torch import nn
+from torchaudio.functional import resample
+from transformers import HubertModel
+
+
+def round_down_nearest_multiple(num, divisor):
+    return num // divisor * divisor
+
+
+def curtail_to_multiple(t, mult, from_left=False):
+    data_len = t.shape[-1]
+    rounded_seq_len = round_down_nearest_multiple(data_len, mult)
+    seq_slice = slice(None, rounded_seq_len) if not from_left else slice(-rounded_seq_len, None)
+    return t[..., seq_slice]
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+class CustomHubert(nn.Module):
+    """
+    checkpoint and kmeans can be downloaded at https://github.com/facebookresearch/fairseq/tree/main/examples/hubert
+    or you can train your own
+    """
+
+    def __init__(self, checkpoint_path, target_sample_hz=16000, seq_len_multiple_of=None, output_layer=9, device=None):
+        super().__init__()
+        self.target_sample_hz = target_sample_hz
+        self.seq_len_multiple_of = seq_len_multiple_of
+        self.output_layer = output_layer
+        if device is not None:
+            self.to(device)
+        self.model = HubertModel.from_pretrained("facebook/hubert-base-ls960")
+        if device is not None:
+            self.model.to(device)
+        self.model.eval()
+
+    @property
+    def groups(self):
+        return 1
+
+    @torch.no_grad()
+    def forward(self, wav_input, flatten=True, input_sample_hz=None):
+        device = wav_input.device
+
+        if exists(input_sample_hz):
+            wav_input = resample(wav_input, input_sample_hz, self.target_sample_hz)
+
+        if exists(self.seq_len_multiple_of):
+            wav_input = curtail_to_multiple(wav_input, self.seq_len_multiple_of)
+
+        outputs = self.model.forward(
+            wav_input,
+            output_hidden_states=True,
+        )
+        embed = outputs["hidden_states"][self.output_layer]
+        embed, packed_shape = pack([embed], "* d")
+        codebook_indices = torch.from_numpy(embed.cpu().detach().numpy()).to(device)
+        if flatten:
+            return codebook_indices
+
+        (codebook_indices,) = unpack(codebook_indices, packed_shape, "*")
+        return codebook_indices

TTS/tts/layers/bark/hubert/tokenizer.py

@@ -0,0 +1,195 @@
+"""
+Custom tokenizer model.
+Author: https://www.github.com/gitmylo/
+License: MIT
+"""
+
+import json
+import os.path
+from zipfile import ZipFile
+
+import numpy
+import torch
+from torch import nn, optim
+
+
+class HubertTokenizer(nn.Module):
+    def __init__(self, hidden_size=1024, input_size=768, output_size=10000, version=0):
+        super().__init__()
+        next_size = input_size
+        if version == 0:
+            self.lstm = nn.LSTM(input_size, hidden_size, 2, batch_first=True)
+            next_size = hidden_size
+        if version == 1:
+            self.lstm = nn.LSTM(input_size, hidden_size, 2, batch_first=True)
+            self.intermediate = nn.Linear(hidden_size, 4096)
+            next_size = 4096
+
+        self.fc = nn.Linear(next_size, output_size)
+        self.softmax = nn.LogSoftmax(dim=1)
+        self.optimizer: optim.Optimizer = None
+        self.lossfunc = nn.CrossEntropyLoss()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.output_size = output_size
+        self.version = version
+
+    def forward(self, x):
+        x, _ = self.lstm(x)
+        if self.version == 1:
+            x = self.intermediate(x)
+        x = self.fc(x)
+        x = self.softmax(x)
+        return x
+
+    @torch.no_grad()
+    def get_token(self, x):
+        """
+        Used to get the token for the first
+        :param x: An array with shape (N, input_size) where N is a whole number greater or equal to 1, and input_size is the input size used when creating the model.
+        :return: An array with shape (N,) where N is the same as N from the input. Every number in the array is a whole number in range 0...output_size - 1 where output_size is the output size used when creating the model.
+        """
+        return torch.argmax(self(x), dim=1)
+
+    def prepare_training(self):
+        self.optimizer = optim.Adam(self.parameters(), 0.001)
+
+    def train_step(self, x_train, y_train, log_loss=False):
+        # y_train = y_train[:-1]
+        # y_train = y_train[1:]
+
+        optimizer = self.optimizer
+        lossfunc = self.lossfunc
+        # Zero the gradients
+        self.zero_grad()
+
+        # Forward pass
+        y_pred = self(x_train)
+
+        y_train_len = len(y_train)
+        y_pred_len = y_pred.shape[0]
+
+        if y_train_len > y_pred_len:
+            diff = y_train_len - y_pred_len
+            y_train = y_train[diff:]
+        elif y_train_len < y_pred_len:
+            diff = y_pred_len - y_train_len
+            y_pred = y_pred[:-diff, :]
+
+        y_train_hot = torch.zeros(len(y_train), self.output_size)
+        y_train_hot[range(len(y_train)), y_train] = 1
+        y_train_hot = y_train_hot.to("cuda")
+
+        # Calculate the loss
+        loss = lossfunc(y_pred, y_train_hot)
+
+        # Print loss
+        if log_loss:
+            print("Loss", loss.item())
+
+        # Backward pass
+        loss.backward()
+
+        # Update the weights
+        optimizer.step()
+
+    def save(self, path):
+        info_path = ".".join(os.path.basename(path).split(".")[:-1]) + "/.info"
+        torch.save(self.state_dict(), path)
+        data_from_model = Data(self.input_size, self.hidden_size, self.output_size, self.version)
+        with ZipFile(path, "a") as model_zip:
+            model_zip.writestr(info_path, data_from_model.save())
+            model_zip.close()
+
+    @staticmethod
+    def load_from_checkpoint(path, map_location=None):
+        old = True
+        with ZipFile(path) as model_zip:
+            filesMatch = [file for file in model_zip.namelist() if file.endswith("/.info")]
+            file = filesMatch[0] if filesMatch else None
+            if file:
+                old = False
+                data_from_model = Data.load(model_zip.read(file).decode("utf-8"))
+            model_zip.close()
+        if old:
+            model = HubertTokenizer()
+        else:
+            model = HubertTokenizer(
+                data_from_model.hidden_size,
+                data_from_model.input_size,
+                data_from_model.output_size,
+                data_from_model.version,
+            )
+        model.load_state_dict(torch.load(path, map_location=map_location))
+        if map_location:
+            model = model.to(map_location)
+        return model
+
+
+class Data:
+    input_size: int
+    hidden_size: int
+    output_size: int
+    version: int
+
+    def __init__(self, input_size=768, hidden_size=1024, output_size=10000, version=0):
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.output_size = output_size
+        self.version = version
+
+    @staticmethod
+    def load(string):
+        data = json.loads(string)
+        return Data(data["input_size"], data["hidden_size"], data["output_size"], data["version"])
+
+    def save(self):
+        data = {
+            "input_size": self.input_size,
+            "hidden_size": self.hidden_size,
+            "output_size": self.output_size,
+            "version": self.version,
+        }
+        return json.dumps(data)
+
+
+def auto_train(data_path, save_path="model.pth", load_model: str = None, save_epochs=1):
+    data_x, data_y = [], []
+
+    if load_model and os.path.isfile(load_model):
+        print("Loading model from", load_model)
+        model_training = HubertTokenizer.load_from_checkpoint(load_model, "cuda")
+    else:
+        print("Creating new model.")
+        model_training = HubertTokenizer(version=1).to("cuda")  # Settings for the model to run without lstm
+    save_path = os.path.join(data_path, save_path)
+    base_save_path = ".".join(save_path.split(".")[:-1])
+
+    sem_string = "_semantic.npy"
+    feat_string = "_semantic_features.npy"
+
+    ready = os.path.join(data_path, "ready")
+    for input_file in os.listdir(ready):
+        full_path = os.path.join(ready, input_file)
+        if input_file.endswith(sem_string):
+            data_y.append(numpy.load(full_path))
+        elif input_file.endswith(feat_string):
+            data_x.append(numpy.load(full_path))
+    model_training.prepare_training()
+
+    epoch = 1
+
+    while 1:
+        for _ in range(save_epochs):
+            j = 0
+            for x, y in zip(data_x, data_y):
+                model_training.train_step(
+                    torch.tensor(x).to("cuda"), torch.tensor(y).to("cuda"), j % 50 == 0
+                )  # Print loss every 50 steps
+                j += 1
+        save_p = save_path
+        save_p_2 = f"{base_save_path}_epoch_{epoch}.pth"
+        model_training.save(save_p)
+        model_training.save(save_p_2)
+        print(f"Epoch {epoch} completed")
+        epoch += 1

TTS/tts/layers/bark/inference_funcs.py

@@ -0,0 +1,606 @@
+import logging
+import os
+import re
+from glob import glob
+from typing import Dict, List
+
+import librosa
+import numpy as np
+import torch
+import torchaudio
+import tqdm
+from encodec.utils import convert_audio
+from scipy.special import softmax
+from torch.nn import functional as F
+
+from TTS.tts.layers.bark.hubert.hubert_manager import HubertManager
+from TTS.tts.layers.bark.hubert.kmeans_hubert import CustomHubert
+from TTS.tts.layers.bark.hubert.tokenizer import HubertTokenizer
+from TTS.tts.layers.bark.load_model import clear_cuda_cache, inference_mode
+
+logger = logging.getLogger(__name__)
+
+
+def _tokenize(tokenizer, text):
+    return tokenizer.encode(text, add_special_tokens=False)
+
+
+def _detokenize(tokenizer, enc_text):
+    return tokenizer.decode(enc_text)
+
+
+def _normalize_whitespace(text):
+    return re.sub(r"\s+", " ", text).strip()
+
+
+def get_voices(extra_voice_dirs: List[str] = []):  # pylint: disable=dangerous-default-value
+    dirs = extra_voice_dirs
+    voices: Dict[str, List[str]] = {}
+    for d in dirs:
+        subs = os.listdir(d)
+        for sub in subs:
+            subj = os.path.join(d, sub)
+            if os.path.isdir(subj):
+                voices[sub] = list(glob(f"{subj}/*.npz"))
+                # fetch audio files if no npz files are found
+                if len(voices[sub]) == 0:
+                    voices[sub] = list(glob(f"{subj}/*.wav")) + list(glob(f"{subj}/*.mp3"))
+    return voices
+
+
+def load_npz(npz_file):
+    x_history = np.load(npz_file)
+    semantic = x_history["semantic_prompt"]
+    coarse = x_history["coarse_prompt"]
+    fine = x_history["fine_prompt"]
+    return semantic, coarse, fine
+
+
+def load_voice(model, voice: str, extra_voice_dirs: List[str] = []):  # pylint: disable=dangerous-default-value
+    if voice == "random":
+        return None, None, None
+
+    voices = get_voices(extra_voice_dirs)
+    paths = voices[voice]
+
+    # bark only uses a single sample for cloning
+    if len(paths) > 1:
+        raise ValueError(f"Voice {voice} has multiple paths: {paths}")
+
+    try:
+        path = voices[voice]
+    except KeyError as e:
+        raise KeyError(f"Voice {voice} not found in {extra_voice_dirs}") from e
+
+    if len(paths) == 1 and paths[0].endswith(".npz"):
+        return load_npz(path[0])
+
+    audio_path = paths[0]
+    # replace the file extension with .npz
+    output_path = os.path.splitext(audio_path)[0] + ".npz"
+    generate_voice(audio=audio_path, model=model, output_path=output_path)
+    return load_voice(model, voice, extra_voice_dirs)
+
+
+def zero_crossing_rate(audio, frame_length=1024, hop_length=512):
+    zero_crossings = np.sum(np.abs(np.diff(np.sign(audio))) / 2)
+    total_frames = 1 + int((len(audio) - frame_length) / hop_length)
+    return zero_crossings / total_frames
+
+
+def compute_spectral_contrast(audio_data, sample_rate, n_bands=6, fmin=200.0):
+    spectral_contrast = librosa.feature.spectral_contrast(y=audio_data, sr=sample_rate, n_bands=n_bands, fmin=fmin)
+    return np.mean(spectral_contrast)
+
+
+def compute_average_bass_energy(audio_data, sample_rate, max_bass_freq=250):
+    stft = librosa.stft(audio_data)
+    power_spectrogram = np.abs(stft) ** 2
+    frequencies = librosa.fft_frequencies(sr=sample_rate, n_fft=stft.shape[0])
+    bass_mask = frequencies <= max_bass_freq
+    bass_energy = power_spectrogram[np.ix_(bass_mask, np.arange(power_spectrogram.shape[1]))].mean()
+    return bass_energy
+
+
+def generate_voice(
+    audio,
+    model,
+    output_path,
+):
+    """Generate a new voice from a given audio and text prompt.
+
+    Args:
+        audio (np.ndarray): The audio to use as a base for the new voice.
+        text (str): Transcription of the audio you are clonning.
+        model (BarkModel): The BarkModel to use for generating the new voice.
+        output_path (str): The path to save the generated voice to.
+    """
+    if isinstance(audio, str):
+        audio, sr = torchaudio.load(audio)
+        audio = convert_audio(audio, sr, model.config.sample_rate, model.encodec.channels)
+        audio = audio.unsqueeze(0).to(model.device)
+
+    with torch.no_grad():
+        encoded_frames = model.encodec.encode(audio)
+    codes = torch.cat([encoded[0] for encoded in encoded_frames], dim=-1).squeeze()  # [n_q, T]
+
+    # move codes to cpu
+    codes = codes.cpu().numpy()
+
+    # generate semantic tokens
+    # Load the HuBERT model
+    hubert_manager = HubertManager()
+    # hubert_manager.make_sure_hubert_installed(model_path=model.config.LOCAL_MODEL_PATHS["hubert"])
+    hubert_manager.make_sure_tokenizer_installed(model_path=model.config.LOCAL_MODEL_PATHS["hubert_tokenizer"])
+
+    hubert_model = CustomHubert(checkpoint_path=model.config.LOCAL_MODEL_PATHS["hubert"]).to(model.device)
+
+    # Load the CustomTokenizer model
+    tokenizer = HubertTokenizer.load_from_checkpoint(
+        model.config.LOCAL_MODEL_PATHS["hubert_tokenizer"], map_location=model.device
+    )
+    # semantic_tokens = model.text_to_semantic(
+    #     text, max_gen_duration_s=seconds, top_k=50, top_p=0.95, temp=0.7
+    # )  # not 100%
+    semantic_vectors = hubert_model.forward(audio[0], input_sample_hz=model.config.sample_rate)
+    semantic_tokens = tokenizer.get_token(semantic_vectors)
+    semantic_tokens = semantic_tokens.cpu().numpy()
+
+    np.savez(output_path, fine_prompt=codes, coarse_prompt=codes[:2, :], semantic_prompt=semantic_tokens)
+
+
+def generate_text_semantic(
+    text,
+    model,
+    history_prompt=None,
+    temp=0.7,
+    top_k=None,
+    top_p=None,
+    silent=False,
+    min_eos_p=0.2,
+    max_gen_duration_s=None,
+    allow_early_stop=True,
+    base=None,
+    use_kv_caching=True,
+    **kwargs,  # pylint: disable=unused-argument
+):
+    """Generate semantic tokens from text.
+
+    Args:
+        text (str): The text to generate semantic tokens from.
+        model (BarkModel): The BarkModel to use for generating the semantic tokens.
+        history_prompt (tuple): A tuple of (semantic_history, coarse_history, fine_history) to use as a prompt for the generation.
+        temp (float): The temperature to use for the generation.
+        top_k (int): The number of top tokens to consider for the generation.
+        top_p (float): The cumulative probability to consider for the generation.
+        silent (bool): Whether to silence the tqdm progress bar.
+        min_eos_p (float): The minimum probability to consider for the end of sentence token.
+        max_gen_duration_s (float): The maximum duration in seconds to generate for.
+        allow_early_stop (bool): Whether to allow the generation to stop early.
+        base (tuple): A tuple of (semantic_history, coarse_history, fine_history) to use as a base for the generation.
+        use_kv_caching (bool): Whether to use key-value caching for the generation.
+        **kwargs: Additional keyword arguments. They are ignored.
+
+    Returns:
+        np.ndarray: The generated semantic tokens.
+    """
+    assert isinstance(text, str)
+    text = _normalize_whitespace(text)
+    assert len(text.strip()) > 0
+    if all(v is not None for v in history_prompt) or base is not None:
+        if history_prompt is not None:
+            semantic_history = history_prompt[0]
+        if base is not None:
+            semantic_history = base[0]
+        assert (
+            isinstance(semantic_history, np.ndarray)
+            and len(semantic_history.shape) == 1
+            and len(semantic_history) > 0
+            and semantic_history.min() >= 0
+            and semantic_history.max() <= model.config.SEMANTIC_VOCAB_SIZE - 1
+        )
+    else:
+        semantic_history = None
+    encoded_text = np.array(_tokenize(model.tokenizer, text)) + model.config.TEXT_ENCODING_OFFSET
+    if len(encoded_text) > 256:
+        p = round((len(encoded_text) - 256) / len(encoded_text) * 100, 1)
+        logger.warning(f"warning, text too long, lopping of last {p}%")
+        encoded_text = encoded_text[:256]
+    encoded_text = np.pad(
+        encoded_text,
+        (0, 256 - len(encoded_text)),
+        constant_values=model.config.TEXT_PAD_TOKEN,
+        mode="constant",
+    )
+    if semantic_history is not None:
+        semantic_history = semantic_history.astype(np.int64)
+        # lop off if history is too long, pad if needed
+        semantic_history = semantic_history[-256:]
+        semantic_history = np.pad(
+            semantic_history,
+            (0, 256 - len(semantic_history)),
+            constant_values=model.config.SEMANTIC_PAD_TOKEN,
+            mode="constant",
+        )
+    else:
+        semantic_history = np.array([model.config.SEMANTIC_PAD_TOKEN] * 256)
+    x = torch.from_numpy(
+        np.hstack([encoded_text, semantic_history, np.array([model.config.SEMANTIC_INFER_TOKEN])]).astype(np.int64)
+    )[None]
+    assert x.shape[1] == 256 + 256 + 1
+    with inference_mode():
+        x = x.to(model.device)
+        n_tot_steps = 768
+        # custom tqdm updates since we don't know when eos will occur
+        pbar = tqdm.tqdm(disable=silent, total=100)
+        pbar_state = 0
+        tot_generated_duration_s = 0
+        kv_cache = None
+        for n in range(n_tot_steps):
+            if use_kv_caching and kv_cache is not None:
+                x_input = x[:, [-1]]
+            else:
+                x_input = x
+            logits, kv_cache = model.semantic_model(
+                x_input, merge_context=True, use_cache=use_kv_caching, past_kv=kv_cache
+            )
+            relevant_logits = logits[0, 0, : model.config.SEMANTIC_VOCAB_SIZE]
+            if allow_early_stop:
+                relevant_logits = torch.hstack(
+                    (relevant_logits, logits[0, 0, [model.config.SEMANTIC_PAD_TOKEN]])
+                )  # eos
+            if top_p is not None:
+                # faster to convert to numpy
+                logits_device = relevant_logits.device
+                logits_dtype = relevant_logits.type()
+                relevant_logits = relevant_logits.detach().cpu().type(torch.float32).numpy()
+                sorted_indices = np.argsort(relevant_logits)[::-1]
+                sorted_logits = relevant_logits[sorted_indices]
+                cumulative_probs = np.cumsum(softmax(sorted_logits))
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].copy()
+                sorted_indices_to_remove[0] = False
+                relevant_logits[sorted_indices[sorted_indices_to_remove]] = -np.inf
+                relevant_logits = torch.from_numpy(relevant_logits)
+                relevant_logits = relevant_logits.to(logits_device).type(logits_dtype)
+            if top_k is not None:
+                v, _ = torch.topk(relevant_logits, min(top_k, relevant_logits.size(-1)))
+                relevant_logits[relevant_logits < v[-1]] = -float("Inf")
+            probs = torch.softmax(relevant_logits / temp, dim=-1)
+            item_next = torch.multinomial(probs, num_samples=1)
+            if allow_early_stop and (
+                item_next == model.config.SEMANTIC_VOCAB_SIZE or (min_eos_p is not None and probs[-1] >= min_eos_p)
+            ):
+                # eos found, so break
+                pbar.update(100 - pbar_state)
+                break
+            x = torch.cat((x, item_next[None]), dim=1)
+            tot_generated_duration_s += 1 / model.config.SEMANTIC_RATE_HZ
+            if max_gen_duration_s is not None and tot_generated_duration_s > max_gen_duration_s:
+                pbar.update(100 - pbar_state)
+                break
+            if n == n_tot_steps - 1:
+                pbar.update(100 - pbar_state)
+                break
+            del logits, relevant_logits, probs, item_next
+            req_pbar_state = np.min([100, int(round(100 * n / n_tot_steps))])
+            if req_pbar_state > pbar_state:
+                pbar.update(req_pbar_state - pbar_state)
+            pbar_state = req_pbar_state
+        pbar.close()
+        out = x.detach().cpu().numpy().squeeze()[256 + 256 + 1 :]
+    assert all(out >= 0) and all(out < model.config.SEMANTIC_VOCAB_SIZE)
+    clear_cuda_cache()
+    return out
+
+
+def _flatten_codebooks(arr, offset_size):
+    assert len(arr.shape) == 2
+    arr = arr.copy()
+    if offset_size is not None:
+        for n in range(1, arr.shape[0]):
+            arr[n, :] += offset_size * n
+    flat_arr = arr.ravel("F")
+    return flat_arr
+
+
+def generate_coarse(
+    x_semantic,
+    model,
+    history_prompt=None,
+    temp=0.7,
+    top_k=None,
+    top_p=None,
+    silent=False,
+    max_coarse_history=630,  # min 60 (faster), max 630 (more context)
+    sliding_window_len=60,
+    base=None,
+    use_kv_caching=True,
+):
+    """Generate coarse audio codes from semantic tokens.
+
+    Args:
+        x_semantic (np.ndarray): The semantic tokens to generate coarse audio codes from.
+        model (BarkModel): The BarkModel to use for generating the coarse audio codes.
+        history_prompt (tuple): A tuple of (semantic_history, coarse_history, fine_history) to use as a prompt for the generation.
+        temp (float): The temperature to use for the generation.
+        top_k (int): The number of top tokens to consider for the generation.
+        top_p (float): The cumulative probability to consider for the generation.
+        silent (bool): Whether to silence the tqdm progress bar.
+        max_coarse_history (int): The maximum number of coarse audio codes to use as history.
+        sliding_window_len (int): The length of the sliding window to use for the generation.
+        base (tuple): A tuple of (semantic_history, coarse_history, fine_history) to use as a base for the generation.
+        use_kv_caching (bool): Whether to use key-value caching for the generation.
+
+    Returns:
+        np.ndarray: The generated coarse audio codes.
+    """
+    assert (
+        isinstance(x_semantic, np.ndarray)
+        and len(x_semantic.shape) == 1
+        and len(x_semantic) > 0
+        and x_semantic.min() >= 0
+        and x_semantic.max() <= model.config.SEMANTIC_VOCAB_SIZE - 1
+    )
+    assert 60 <= max_coarse_history <= 630
+    assert max_coarse_history + sliding_window_len <= 1024 - 256
+    semantic_to_coarse_ratio = (
+        model.config.COARSE_RATE_HZ / model.config.SEMANTIC_RATE_HZ * model.config.N_COARSE_CODEBOOKS
+    )
+    max_semantic_history = int(np.floor(max_coarse_history / semantic_to_coarse_ratio))
+    if all(v is not None for v in history_prompt) or base is not None:
+        if history_prompt is not None:
+            x_history = history_prompt
+            x_semantic_history = x_history[0]
+            x_coarse_history = x_history[1]
+        if base is not None:
+            x_semantic_history = base[0]
+            x_coarse_history = base[1]
+        assert (
+            isinstance(x_semantic_history, np.ndarray)
+            and len(x_semantic_history.shape) == 1
+            and len(x_semantic_history) > 0
+            and x_semantic_history.min() >= 0
+            and x_semantic_history.max() <= model.config.SEMANTIC_VOCAB_SIZE - 1
+            and isinstance(x_coarse_history, np.ndarray)
+            and len(x_coarse_history.shape) == 2
+            and x_coarse_history.shape[0] == model.config.N_COARSE_CODEBOOKS
+            and x_coarse_history.shape[-1] >= 0
+            and x_coarse_history.min() >= 0
+            and x_coarse_history.max() <= model.config.CODEBOOK_SIZE - 1
+            and (
+                round(x_coarse_history.shape[-1] / len(x_semantic_history), 1)
+                == round(semantic_to_coarse_ratio / model.config.N_COARSE_CODEBOOKS, 1)
+            )
+        )
+        x_coarse_history = (
+            _flatten_codebooks(x_coarse_history, model.config.CODEBOOK_SIZE) + model.config.SEMANTIC_VOCAB_SIZE
+        )
+        # trim histories correctly
+        n_semantic_hist_provided = np.min(
+            [
+                max_semantic_history,
+                len(x_semantic_history) - len(x_semantic_history) % 2,
+                int(np.floor(len(x_coarse_history) / semantic_to_coarse_ratio)),
+            ]
+        )
+        n_coarse_hist_provided = int(round(n_semantic_hist_provided * semantic_to_coarse_ratio))
+        x_semantic_history = x_semantic_history[-n_semantic_hist_provided:].astype(np.int32)
+        x_coarse_history = x_coarse_history[-n_coarse_hist_provided:].astype(np.int32)
+        # TODO: bit of a hack for time alignment (sounds better)
+        x_coarse_history = x_coarse_history[:-2]
+    else:
+        x_semantic_history = np.array([], dtype=np.int32)
+        x_coarse_history = np.array([], dtype=np.int32)
+    # start loop
+    n_steps = int(
+        round(
+            np.floor(len(x_semantic) * semantic_to_coarse_ratio / model.config.N_COARSE_CODEBOOKS)
+            * model.config.N_COARSE_CODEBOOKS
+        )
+    )
+    assert n_steps > 0 and n_steps % model.config.N_COARSE_CODEBOOKS == 0
+    x_semantic = np.hstack([x_semantic_history, x_semantic]).astype(np.int32)
+    x_coarse = x_coarse_history.astype(np.int32)
+    base_semantic_idx = len(x_semantic_history)
+    with inference_mode():
+        x_semantic_in = torch.from_numpy(x_semantic)[None].to(model.device)
+        x_coarse_in = torch.from_numpy(x_coarse)[None].to(model.device)
+        n_window_steps = int(np.ceil(n_steps / sliding_window_len))
+        n_step = 0
+        for _ in tqdm.tqdm(range(n_window_steps), total=n_window_steps, disable=silent):
+            semantic_idx = base_semantic_idx + int(round(n_step / semantic_to_coarse_ratio))
+            # pad from right side
+            x_in = x_semantic_in[:, np.max([0, semantic_idx - max_semantic_history]) :]
+            x_in = x_in[:, :256]
+            x_in = F.pad(
+                x_in,
+                (0, 256 - x_in.shape[-1]),
+                "constant",
+                model.config.COARSE_SEMANTIC_PAD_TOKEN,
+            )
+            x_in = torch.hstack(
+                [
+                    x_in,
+                    torch.tensor([model.config.COARSE_INFER_TOKEN])[None].to(model.device),
+                    x_coarse_in[:, -max_coarse_history:],
+                ]
+            )
+            kv_cache = None
+            for _ in range(sliding_window_len):
+                if n_step >= n_steps:
+                    continue
+                is_major_step = n_step % model.config.N_COARSE_CODEBOOKS == 0
+
+                if use_kv_caching and kv_cache is not None:
+                    x_input = x_in[:, [-1]]
+                else:
+                    x_input = x_in
+
+                logits, kv_cache = model.coarse_model(x_input, use_cache=use_kv_caching, past_kv=kv_cache)
+                logit_start_idx = (
+                    model.config.SEMANTIC_VOCAB_SIZE + (1 - int(is_major_step)) * model.config.CODEBOOK_SIZE
+                )
+                logit_end_idx = model.config.SEMANTIC_VOCAB_SIZE + (2 - int(is_major_step)) * model.config.CODEBOOK_SIZE
+                relevant_logits = logits[0, 0, logit_start_idx:logit_end_idx]
+                if top_p is not None:
+                    # faster to convert to numpy
+                    logits_device = relevant_logits.device
+                    logits_dtype = relevant_logits.type()
+                    relevant_logits = relevant_logits.detach().cpu().type(torch.float32).numpy()
+                    sorted_indices = np.argsort(relevant_logits)[::-1]
+                    sorted_logits = relevant_logits[sorted_indices]
+                    cumulative_probs = np.cumsum(torch.nn.functional.softmax(sorted_logits))
+                    sorted_indices_to_remove = cumulative_probs > top_p
+                    sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].copy()
+                    sorted_indices_to_remove[0] = False
+                    relevant_logits[sorted_indices[sorted_indices_to_remove]] = -np.inf
+                    relevant_logits = torch.from_numpy(relevant_logits)
+                    relevant_logits = relevant_logits.to(logits_device).type(logits_dtype)
+                if top_k is not None:
+                    v, _ = torch.topk(relevant_logits, min(top_k, relevant_logits.size(-1)))
+                    relevant_logits[relevant_logits < v[-1]] = -float("Inf")
+                probs = torch.nn.functional.softmax(relevant_logits / temp, dim=-1)
+                item_next = torch.multinomial(probs, num_samples=1)
+                item_next += logit_start_idx
+                x_coarse_in = torch.cat((x_coarse_in, item_next[None]), dim=1)
+                x_in = torch.cat((x_in, item_next[None]), dim=1)
+                del logits, relevant_logits, probs, item_next
+                n_step += 1
+            del x_in
+        del x_semantic_in
+    gen_coarse_arr = x_coarse_in.detach().cpu().numpy().squeeze()[len(x_coarse_history) :]
+    del x_coarse_in
+    assert len(gen_coarse_arr) == n_steps
+    gen_coarse_audio_arr = (
+        gen_coarse_arr.reshape(-1, model.config.N_COARSE_CODEBOOKS).T - model.config.SEMANTIC_VOCAB_SIZE
+    )
+    for n in range(1, model.config.N_COARSE_CODEBOOKS):
+        gen_coarse_audio_arr[n, :] -= n * model.config.CODEBOOK_SIZE
+    clear_cuda_cache()
+    return gen_coarse_audio_arr
+
+
+def generate_fine(
+    x_coarse_gen,
+    model,
+    history_prompt=None,
+    temp=0.5,
+    silent=True,
+    base=None,
+):
+    """Generate full audio codes from coarse audio codes.
+
+    Args:
+        x_coarse_gen (np.ndarray): The coarse audio codes to generate full audio codes from.
+        model (BarkModel): The BarkModel to use for generating the full audio codes.
+        history_prompt (tuple): A tuple of (semantic_history, coarse_history, fine_history) to use as a prompt for the generation.
+        temp (float): The temperature to use for the generation.
+        silent (bool): Whether to silence the tqdm progress bar.
+        base (tuple): A tuple of (semantic_history, coarse_history, fine_history) to use as a base for the generation.
+
+    Returns:
+        np.ndarray: The generated full audio codes.
+    """
+    assert (
+        isinstance(x_coarse_gen, np.ndarray)
+        and len(x_coarse_gen.shape) == 2
+        and 1 <= x_coarse_gen.shape[0] <= model.config.N_FINE_CODEBOOKS - 1
+        and x_coarse_gen.shape[1] > 0
+        and x_coarse_gen.min() >= 0
+        and x_coarse_gen.max() <= model.config.CODEBOOK_SIZE - 1
+    )
+    if all(v is not None for v in history_prompt) or base is not None:
+        if history_prompt is not None:
+            x_fine_history = history_prompt[2]
+        if base is not None:
+            x_fine_history = base[2]
+        assert (
+            isinstance(x_fine_history, np.ndarray)
+            and len(x_fine_history.shape) == 2
+            and x_fine_history.shape[0] == model.config.N_FINE_CODEBOOKS
+            and x_fine_history.shape[1] >= 0
+            and x_fine_history.min() >= 0
+            and x_fine_history.max() <= model.config.CODEBOOK_SIZE - 1
+        )
+    else:
+        x_fine_history = None
+    n_coarse = x_coarse_gen.shape[0]
+    # make input arr
+    in_arr = np.vstack(
+        [
+            x_coarse_gen,
+            np.zeros((model.config.N_FINE_CODEBOOKS - n_coarse, x_coarse_gen.shape[1]))
+            + model.config.CODEBOOK_SIZE,  # padding
+        ]
+    ).astype(np.int32)
+    # prepend history if available (max 512)
+    if x_fine_history is not None:
+        x_fine_history = x_fine_history.astype(np.int32)
+        in_arr = np.hstack(
+            [
+                x_fine_history[:, -512:].astype(np.int32),
+                in_arr,
+            ]
+        )
+        n_history = x_fine_history[:, -512:].shape[1]
+    else:
+        n_history = 0
+    n_remove_from_end = 0
+    # need to pad if too short (since non-causal model)
+    if in_arr.shape[1] < 1024:
+        n_remove_from_end = 1024 - in_arr.shape[1]
+        in_arr = np.hstack(
+            [
+                in_arr,
+                np.zeros((model.config.N_FINE_CODEBOOKS, n_remove_from_end), dtype=np.int32)
+                + model.config.CODEBOOK_SIZE,
+            ]
+        )
+    # we can be lazy about fractional loop and just keep overwriting codebooks
+    n_loops = np.max([0, int(np.ceil((x_coarse_gen.shape[1] - (1024 - n_history)) / 512))]) + 1
+    with inference_mode():
+        in_arr = torch.tensor(in_arr.T).to(model.device)
+        for n in tqdm.tqdm(range(n_loops), disable=silent):
+            start_idx = np.min([n * 512, in_arr.shape[0] - 1024])
+            start_fill_idx = np.min([n_history + n * 512, in_arr.shape[0] - 512])
+            rel_start_fill_idx = start_fill_idx - start_idx
+            in_buffer = in_arr[start_idx : start_idx + 1024, :][None]
+            for nn in range(n_coarse, model.config.N_FINE_CODEBOOKS):
+                logits = model.fine_model(nn, in_buffer)
+                if temp is None:
+                    relevant_logits = logits[0, rel_start_fill_idx:, : model.config.CODEBOOK_SIZE]
+                    codebook_preds = torch.argmax(relevant_logits, -1)
+                else:
+                    relevant_logits = logits[0, :, : model.config.CODEBOOK_SIZE] / temp
+                    probs = F.softmax(relevant_logits, dim=-1)
+                    codebook_preds = torch.hstack(
+                        [torch.multinomial(probs[n], num_samples=1) for n in range(rel_start_fill_idx, 1024)]
+                    )
+                in_buffer[0, rel_start_fill_idx:, nn] = codebook_preds
+                del logits, codebook_preds
+            # transfer over info into model_in and convert to numpy
+            for nn in range(n_coarse, model.config.N_FINE_CODEBOOKS):
+                in_arr[start_fill_idx : start_fill_idx + (1024 - rel_start_fill_idx), nn] = in_buffer[
+                    0, rel_start_fill_idx:, nn
+                ]
+            del in_buffer
+        gen_fine_arr = in_arr.detach().cpu().numpy().squeeze().T
+        del in_arr
+    gen_fine_arr = gen_fine_arr[:, n_history:]
+    if n_remove_from_end > 0:
+        gen_fine_arr = gen_fine_arr[:, :-n_remove_from_end]
+    assert gen_fine_arr.shape[-1] == x_coarse_gen.shape[-1]
+    clear_cuda_cache()
+    return gen_fine_arr
+
+
+def codec_decode(fine_tokens, model):
+    """Turn quantized audio codes into audio array using encodec."""
+    arr = torch.from_numpy(fine_tokens)[None]
+    arr = arr.to(model.device)
+    arr = arr.transpose(0, 1)
+    emb = model.encodec.quantizer.decode(arr)
+    out = model.encodec.decoder(emb)
+    audio_arr = out.detach().cpu().numpy().squeeze()
+    return audio_arr

TTS/tts/layers/bark/load_model.py

@@ -0,0 +1,160 @@
+import contextlib
+import functools
+import hashlib
+import logging
+import os
+
+import requests
+import torch
+import tqdm
+
+from TTS.tts.layers.bark.model import GPT, GPTConfig
+from TTS.tts.layers.bark.model_fine import FineGPT, FineGPTConfig
+
+if (
+    torch.cuda.is_available()
+    and hasattr(torch.cuda, "amp")
+    and hasattr(torch.cuda.amp, "autocast")
+    and torch.cuda.is_bf16_supported()
+):
+    autocast = functools.partial(torch.cuda.amp.autocast, dtype=torch.bfloat16)
+else:
+
+    @contextlib.contextmanager
+    def autocast():
+        yield
+
+
+# hold models in global scope to lazy load
+
+logger = logging.getLogger(__name__)
+
+
+if not hasattr(torch.nn.functional, "scaled_dot_product_attention"):
+    logger.warning(
+        "torch version does not support flash attention. You will get significantly faster"
+        + " inference speed by upgrade torch to newest version / nightly."
+    )
+
+
+def _md5(fname):
+    hash_md5 = hashlib.md5()
+    with open(fname, "rb") as f:
+        for chunk in iter(lambda: f.read(4096), b""):
+            hash_md5.update(chunk)
+    return hash_md5.hexdigest()
+
+
+def _download(from_s3_path, to_local_path, CACHE_DIR):
+    os.makedirs(CACHE_DIR, exist_ok=True)
+    response = requests.get(from_s3_path, stream=True)
+    total_size_in_bytes = int(response.headers.get("content-length", 0))
+    block_size = 1024  # 1 Kibibyte
+    progress_bar = tqdm.tqdm(total=total_size_in_bytes, unit="iB", unit_scale=True)
+    with open(to_local_path, "wb") as file:
+        for data in response.iter_content(block_size):
+            progress_bar.update(len(data))
+            file.write(data)
+    progress_bar.close()
+    if total_size_in_bytes not in [0, progress_bar.n]:
+        raise ValueError("ERROR, something went wrong")
+
+
+class InferenceContext:
+    def __init__(self, benchmark=False):
+        # we can't expect inputs to be the same length, so disable benchmarking by default
+        self._chosen_cudnn_benchmark = benchmark
+        self._cudnn_benchmark = None
+
+    def __enter__(self):
+        self._cudnn_benchmark = torch.backends.cudnn.benchmark
+        torch.backends.cudnn.benchmark = self._chosen_cudnn_benchmark
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        torch.backends.cudnn.benchmark = self._cudnn_benchmark
+
+
+if torch.cuda.is_available():
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+
+@contextlib.contextmanager
+def inference_mode():
+    with InferenceContext(), torch.inference_mode(), torch.no_grad(), autocast():
+        yield
+
+
+def clear_cuda_cache():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+
+def load_model(ckpt_path, device, config, model_type="text"):
+    logger.info(f"loading {model_type} model from {ckpt_path}...")
+
+    if device == "cpu":
+        logger.warning("No GPU being used. Careful, Inference might be extremely slow!")
+    if model_type == "text":
+        ConfigClass = GPTConfig
+        ModelClass = GPT
+    elif model_type == "coarse":
+        ConfigClass = GPTConfig
+        ModelClass = GPT
+    elif model_type == "fine":
+        ConfigClass = FineGPTConfig
+        ModelClass = FineGPT
+    else:
+        raise NotImplementedError()
+    if (
+        not config.USE_SMALLER_MODELS
+        and os.path.exists(ckpt_path)
+        and _md5(ckpt_path) != config.REMOTE_MODEL_PATHS[model_type]["checksum"]
+    ):
+        logger.warning(f"found outdated {model_type} model, removing...")
+        os.remove(ckpt_path)
+    if not os.path.exists(ckpt_path):
+        logger.info(f"{model_type} model not found, downloading...")
+        _download(config.REMOTE_MODEL_PATHS[model_type]["path"], ckpt_path, config.CACHE_DIR)
+
+    checkpoint = torch.load(ckpt_path, map_location=device)
+    # this is a hack
+    model_args = checkpoint["model_args"]
+    if "input_vocab_size" not in model_args:
+        model_args["input_vocab_size"] = model_args["vocab_size"]
+        model_args["output_vocab_size"] = model_args["vocab_size"]
+        del model_args["vocab_size"]
+
+    gptconf = ConfigClass(**checkpoint["model_args"])
+    if model_type == "text":
+        config.semantic_config = gptconf
+    elif model_type == "coarse":
+        config.coarse_config = gptconf
+    elif model_type == "fine":
+        config.fine_config = gptconf
+
+    model = ModelClass(gptconf)
+    state_dict = checkpoint["model"]
+    # fixup checkpoint
+    unwanted_prefix = "_orig_mod."
+    for k, _ in list(state_dict.items()):
+        if k.startswith(unwanted_prefix):
+            state_dict[k[len(unwanted_prefix) :]] = state_dict.pop(k)
+    extra_keys = set(state_dict.keys()) - set(model.state_dict().keys())
+    extra_keys = set(k for k in extra_keys if not k.endswith(".attn.bias"))
+    missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())
+    missing_keys = set(k for k in missing_keys if not k.endswith(".attn.bias"))
+    if len(extra_keys) != 0:
+        raise ValueError(f"extra keys found: {extra_keys}")
+    if len(missing_keys) != 0:
+        raise ValueError(f"missing keys: {missing_keys}")
+    model.load_state_dict(state_dict, strict=False)
+    n_params = model.get_num_params()
+    val_loss = checkpoint["best_val_loss"].item()
+    logger.info(f"model loaded: {round(n_params/1e6,1)}M params, {round(val_loss,3)} loss")
+    model.eval()
+    model.to(device)
+    del checkpoint, state_dict
+    clear_cuda_cache()
+    return model, config

TTS/tts/layers/bark/model.py

@@ -0,0 +1,233 @@
+"""
+Much of this code is adapted from Andrej Karpathy's NanoGPT
+(https://github.com/karpathy/nanoGPT)
+"""
+import math
+from dataclasses import dataclass
+
+import torch
+from coqpit import Coqpit
+from torch import nn
+from torch.nn import functional as F
+
+
+class LayerNorm(nn.Module):
+    """LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False"""
+
+    def __init__(self, ndim, bias):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(ndim))
+        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
+
+    def forward(self, x):
+        return F.layer_norm(x, self.weight.shape, self.weight, self.bias, 1e-5)
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        # key, query, value projections for all heads, but in a batch
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        # output projection
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        # regularization
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.dropout = config.dropout
+        # flash attention make GPU go brrrrr but support is only in PyTorch nightly and still a bit scary
+        self.flash = hasattr(torch.nn.functional, "scaled_dot_product_attention")
+        if not self.flash:
+            # print("WARNING: using slow attention. Flash Attention atm needs PyTorch nightly and dropout=0.0")
+            # causal mask to ensure that attention is only applied to the left in the input sequence
+            self.register_buffer(
+                "bias",
+                torch.tril(torch.ones(config.block_size, config.block_size)).view(
+                    1, 1, config.block_size, config.block_size
+                ),
+            )
+
+    def forward(self, x, past_kv=None, use_cache=False):
+        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)
+
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
+
+        if past_kv is not None:
+            past_key = past_kv[0]
+            past_value = past_kv[1]
+            k = torch.cat((past_key, k), dim=-2)
+            v = torch.cat((past_value, v), dim=-2)
+
+        FULL_T = k.shape[-2]
+
+        if use_cache is True:
+            present = (k, v)
+        else:
+            present = None
+
+        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        if self.flash:
+            # efficient attention using Flash Attention CUDA kernels
+            if past_kv is not None:
+                # When `past_kv` is provided, we're doing incremental decoding and `q.shape[2] == 1`: q only contains
+                # the query for the last token. scaled_dot_product_attention interprets this as the first token in the
+                # sequence, so if is_causal=True it will mask out all attention from it. This is not what we want, so
+                # to work around this we set is_causal=False.
+                is_causal = False
+            else:
+                is_causal = True
+
+            # efficient attention using Flash Attention CUDA kernels
+            y = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=self.dropout, is_causal=is_causal)
+        else:
+            # manual implementation of attention
+            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+            att = att.masked_fill(self.bias[:, :, FULL_T - T : FULL_T, :FULL_T] == 0, float("-inf"))
+            att = F.softmax(att, dim=-1)
+            att = self.attn_dropout(att)
+            y = att @ v  # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+        y = y.transpose(1, 2).contiguous().view(B, T, C)  # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.resid_dropout(self.c_proj(y))
+        return (y, present)
+
+
+class MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+        self.gelu = nn.GELU()
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_proj(x)
+        x = self.dropout(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
+        self.mlp = MLP(config)
+        self.layer_idx = layer_idx
+
+    def forward(self, x, past_kv=None, use_cache=False):
+        attn_output, prev_kvs = self.attn(self.ln_1(x), past_kv=past_kv, use_cache=use_cache)
+        x = x + attn_output
+        x = x + self.mlp(self.ln_2(x))
+        return (x, prev_kvs)
+
+
+@dataclass
+class GPTConfig(Coqpit):
+    block_size: int = 1024
+    input_vocab_size: int = 10_048
+    output_vocab_size: int = 10_048
+    n_layer: int = 12
+    n_head: int = 12
+    n_embd: int = 768
+    dropout: float = 0.0
+    bias: bool = True  # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster
+
+
+class GPT(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.input_vocab_size is not None
+        assert config.output_vocab_size is not None
+        assert config.block_size is not None
+        self.config = config
+
+        self.transformer = nn.ModuleDict(
+            dict(
+                wte=nn.Embedding(config.input_vocab_size, config.n_embd),
+                wpe=nn.Embedding(config.block_size, config.n_embd),
+                drop=nn.Dropout(config.dropout),
+                h=nn.ModuleList([Block(config, idx) for idx in range(config.n_layer)]),
+                ln_f=LayerNorm(config.n_embd, bias=config.bias),
+            )
+        )
+        self.lm_head = nn.Linear(config.n_embd, config.output_vocab_size, bias=False)
+
+    def get_num_params(self, non_embedding=True):
+        """
+        Return the number of parameters in the model.
+        For non-embedding count (default), the position embeddings get subtracted.
+        The token embeddings would too, except due to the parameter sharing these
+        params are actually used as weights in the final layer, so we include them.
+        """
+        n_params = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            n_params -= self.transformer.wte.weight.numel()
+            n_params -= self.transformer.wpe.weight.numel()
+        return n_params
+
+    def forward(self, idx, merge_context=False, past_kv=None, position_ids=None, use_cache=False):
+        device = idx.device
+        _, t = idx.size()
+        if past_kv is not None:
+            assert t == 1
+            tok_emb = self.transformer.wte(idx)  # token embeddings of shape (b, t, n_embd)
+        else:
+            if merge_context:
+                assert idx.shape[1] >= 256 + 256 + 1
+                t = idx.shape[1] - 256
+            else:
+                assert (
+                    t <= self.config.block_size
+                ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+
+            # forward the GPT model itself
+            if merge_context:
+                tok_emb = torch.cat(
+                    [
+                        self.transformer.wte(idx[:, :256]) + self.transformer.wte(idx[:, 256 : 256 + 256]),
+                        self.transformer.wte(idx[:, 256 + 256 :]),
+                    ],
+                    dim=1,
+                )
+            else:
+                tok_emb = self.transformer.wte(idx)  # token embeddings of shape (b, t, n_embd)
+
+        if past_kv is None:
+            past_length = 0
+            past_kv = tuple([None] * len(self.transformer.h))
+        else:
+            past_length = past_kv[0][0].size(-2)
+
+        if position_ids is None:
+            position_ids = torch.arange(past_length, t + past_length, dtype=torch.long, device=device)
+            position_ids = position_ids.unsqueeze(0)  # shape (1, t)
+            assert position_ids.shape == (1, t)
+
+        pos_emb = self.transformer.wpe(position_ids)  # position embeddings of shape (1, t, n_embd)
+
+        x = self.transformer.drop(tok_emb + pos_emb)
+
+        new_kv = () if use_cache else None
+
+        for _, (block, past_layer_kv) in enumerate(zip(self.transformer.h, past_kv)):
+            x, kv = block(x, past_kv=past_layer_kv, use_cache=use_cache)
+
+            if use_cache:
+                new_kv = new_kv + (kv,)
+
+        x = self.transformer.ln_f(x)
+
+        # inference-time mini-optimization: only forward the lm_head on the very last position
+        logits = self.lm_head(x[:, [-1], :])  # note: using list [-1] to preserve the time dim
+
+        return (logits, new_kv)

TTS/tts/layers/bark/model_fine.py

@@ -0,0 +1,142 @@
+"""
+Much of this code is adapted from Andrej Karpathy's NanoGPT
+(https://github.com/karpathy/nanoGPT)
+"""
+import math
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from .model import GPT, MLP, GPTConfig
+
+
+class NonCausalSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        # key, query, value projections for all heads, but in a batch
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        # output projection
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        # regularization
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.dropout = config.dropout
+        # flash attention make GPU go brrrrr but support is only in PyTorch nightly and still a bit scary
+        self.flash = hasattr(torch.nn.functional, "scaled_dot_product_attention") and self.dropout == 0.0
+
+    def forward(self, x):
+        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)
+
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
+
+        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        if self.flash:
+            # efficient attention using Flash Attention CUDA kernels
+            y = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, attn_mask=None, dropout_p=self.dropout, is_causal=False
+            )
+        else:
+            # manual implementation of attention
+            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+            att = F.softmax(att, dim=-1)
+            att = self.attn_dropout(att)
+            y = att @ v  # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+        y = y.transpose(1, 2).contiguous().view(B, T, C)  # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+
+class FineBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(config.n_embd)
+        self.attn = NonCausalSelfAttention(config)
+        self.ln_2 = nn.LayerNorm(config.n_embd)
+        self.mlp = MLP(config)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class FineGPT(GPT):
+    def __init__(self, config):
+        super().__init__(config)
+        del self.lm_head
+        self.config = config
+        self.n_codes_total = config.n_codes_total
+        self.transformer = nn.ModuleDict(
+            dict(
+                wtes=nn.ModuleList(
+                    [nn.Embedding(config.input_vocab_size, config.n_embd) for _ in range(config.n_codes_total)]
+                ),
+                wpe=nn.Embedding(config.block_size, config.n_embd),
+                drop=nn.Dropout(config.dropout),
+                h=nn.ModuleList([FineBlock(config) for _ in range(config.n_layer)]),
+                ln_f=nn.LayerNorm(config.n_embd),
+            )
+        )
+        self.lm_heads = nn.ModuleList(
+            [
+                nn.Linear(config.n_embd, config.output_vocab_size, bias=False)
+                for _ in range(config.n_codes_given, self.n_codes_total)
+            ]
+        )
+        for i in range(self.n_codes_total - config.n_codes_given):
+            self.transformer.wtes[i + 1].weight = self.lm_heads[i].weight
+
+    def forward(self, pred_idx, idx):
+        device = idx.device
+        b, t, codes = idx.size()
+        assert (
+            t <= self.config.block_size
+        ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+        assert pred_idx > 0, "cannot predict 0th codebook"
+        assert codes == self.n_codes_total, (b, t, codes)
+        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0)  # shape (1, t)
+
+        # forward the GPT model itself
+        tok_embs = [
+            wte(idx[:, :, i]).unsqueeze(-1) for i, wte in enumerate(self.transformer.wtes)
+        ]  # token embeddings of shape (b, t, n_embd)
+        tok_emb = torch.cat(tok_embs, dim=-1)
+        pos_emb = self.transformer.wpe(pos)  # position embeddings of shape (1, t, n_embd)
+        x = tok_emb[:, :, :, : pred_idx + 1].sum(dim=-1)
+        x = self.transformer.drop(x + pos_emb)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+        logits = self.lm_heads[pred_idx - self.config.n_codes_given](x)
+        return logits
+
+    def get_num_params(self, non_embedding=True):
+        """
+        Return the number of parameters in the model.
+        For non-embedding count (default), the position embeddings get subtracted.
+        The token embeddings would too, except due to the parameter sharing these
+        params are actually used as weights in the final layer, so we include them.
+        """
+        n_params = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            for wte in self.transformer.wtes:
+                n_params -= wte.weight.numel()
+            n_params -= self.transformer.wpe.weight.numel()
+        return n_params
+
+
+@dataclass
+class FineGPTConfig(GPTConfig):
+    n_codes_total: int = 8
+    n_codes_given: int = 1

TTS/tts/layers/delightful_tts/acoustic_model.py

@@ -0,0 +1,563 @@
+### credit: https://github.com/dunky11/voicesmith
+from typing import Callable, Dict, Tuple
+
+import torch
+import torch.nn.functional as F
+from coqpit import Coqpit
+from torch import nn
+
+from TTS.tts.layers.delightful_tts.conformer import Conformer
+from TTS.tts.layers.delightful_tts.encoders import (
+    PhonemeLevelProsodyEncoder,
+    UtteranceLevelProsodyEncoder,
+    get_mask_from_lengths,
+)
+from TTS.tts.layers.delightful_tts.energy_adaptor import EnergyAdaptor
+from TTS.tts.layers.delightful_tts.networks import EmbeddingPadded, positional_encoding
+from TTS.tts.layers.delightful_tts.phoneme_prosody_predictor import PhonemeProsodyPredictor
+from TTS.tts.layers.delightful_tts.pitch_adaptor import PitchAdaptor
+from TTS.tts.layers.delightful_tts.variance_predictor import VariancePredictor
+from TTS.tts.layers.generic.aligner import AlignmentNetwork
+from TTS.tts.utils.helpers import generate_path, maximum_path, sequence_mask
+
+
+class AcousticModel(torch.nn.Module):
+    def __init__(
+        self,
+        args: "ModelArgs",
+        tokenizer: "TTSTokenizer" = None,
+        speaker_manager: "SpeakerManager" = None,
+    ):
+        super().__init__()
+        self.args = args
+        self.tokenizer = tokenizer
+        self.speaker_manager = speaker_manager
+
+        self.init_multispeaker(args)
+        # self.set_embedding_dims()
+
+        self.length_scale = (
+            float(self.args.length_scale) if isinstance(self.args.length_scale, int) else self.args.length_scale
+        )
+
+        self.emb_dim = args.n_hidden_conformer_encoder
+        self.encoder = Conformer(
+            dim=self.args.n_hidden_conformer_encoder,
+            n_layers=self.args.n_layers_conformer_encoder,
+            n_heads=self.args.n_heads_conformer_encoder,
+            speaker_embedding_dim=self.embedded_speaker_dim,
+            p_dropout=self.args.dropout_conformer_encoder,
+            kernel_size_conv_mod=self.args.kernel_size_conv_mod_conformer_encoder,
+            lrelu_slope=self.args.lrelu_slope,
+        )
+        self.pitch_adaptor = PitchAdaptor(
+            n_input=self.args.n_hidden_conformer_encoder,
+            n_hidden=self.args.n_hidden_variance_adaptor,
+            n_out=1,
+            kernel_size=self.args.kernel_size_variance_adaptor,
+            emb_kernel_size=self.args.emb_kernel_size_variance_adaptor,
+            p_dropout=self.args.dropout_variance_adaptor,
+            lrelu_slope=self.args.lrelu_slope,
+        )
+        self.energy_adaptor = EnergyAdaptor(
+            channels_in=self.args.n_hidden_conformer_encoder,
+            channels_hidden=self.args.n_hidden_variance_adaptor,
+            channels_out=1,
+            kernel_size=self.args.kernel_size_variance_adaptor,
+            emb_kernel_size=self.args.emb_kernel_size_variance_adaptor,
+            dropout=self.args.dropout_variance_adaptor,
+            lrelu_slope=self.args.lrelu_slope,
+        )
+
+        self.aligner = AlignmentNetwork(
+            in_query_channels=self.args.out_channels,
+            in_key_channels=self.args.n_hidden_conformer_encoder,
+        )
+
+        self.duration_predictor = VariancePredictor(
+            channels_in=self.args.n_hidden_conformer_encoder,
+            channels=self.args.n_hidden_variance_adaptor,
+            channels_out=1,
+            kernel_size=self.args.kernel_size_variance_adaptor,
+            p_dropout=self.args.dropout_variance_adaptor,
+            lrelu_slope=self.args.lrelu_slope,
+        )
+
+        self.utterance_prosody_encoder = UtteranceLevelProsodyEncoder(
+            num_mels=self.args.num_mels,
+            ref_enc_filters=self.args.ref_enc_filters_reference_encoder,
+            ref_enc_size=self.args.ref_enc_size_reference_encoder,
+            ref_enc_gru_size=self.args.ref_enc_gru_size_reference_encoder,
+            ref_enc_strides=self.args.ref_enc_strides_reference_encoder,
+            n_hidden=self.args.n_hidden_conformer_encoder,
+            dropout=self.args.dropout_conformer_encoder,
+            bottleneck_size_u=self.args.bottleneck_size_u_reference_encoder,
+            token_num=self.args.token_num_reference_encoder,
+        )
+
+        self.utterance_prosody_predictor = PhonemeProsodyPredictor(
+            hidden_size=self.args.n_hidden_conformer_encoder,
+            kernel_size=self.args.predictor_kernel_size_reference_encoder,
+            dropout=self.args.dropout_conformer_encoder,
+            bottleneck_size=self.args.bottleneck_size_u_reference_encoder,
+            lrelu_slope=self.args.lrelu_slope,
+        )
+
+        self.phoneme_prosody_encoder = PhonemeLevelProsodyEncoder(
+            num_mels=self.args.num_mels,
+            ref_enc_filters=self.args.ref_enc_filters_reference_encoder,
+            ref_enc_size=self.args.ref_enc_size_reference_encoder,
+            ref_enc_gru_size=self.args.ref_enc_gru_size_reference_encoder,
+            ref_enc_strides=self.args.ref_enc_strides_reference_encoder,
+            n_hidden=self.args.n_hidden_conformer_encoder,
+            dropout=self.args.dropout_conformer_encoder,
+            bottleneck_size_p=self.args.bottleneck_size_p_reference_encoder,
+            n_heads=self.args.n_heads_conformer_encoder,
+        )
+
+        self.phoneme_prosody_predictor = PhonemeProsodyPredictor(
+            hidden_size=self.args.n_hidden_conformer_encoder,
+            kernel_size=self.args.predictor_kernel_size_reference_encoder,
+            dropout=self.args.dropout_conformer_encoder,
+            bottleneck_size=self.args.bottleneck_size_p_reference_encoder,
+            lrelu_slope=self.args.lrelu_slope,
+        )
+
+        self.u_bottle_out = nn.Linear(
+            self.args.bottleneck_size_u_reference_encoder,
+            self.args.n_hidden_conformer_encoder,
+        )
+
+        self.u_norm = nn.InstanceNorm1d(self.args.bottleneck_size_u_reference_encoder)
+        self.p_bottle_out = nn.Linear(
+            self.args.bottleneck_size_p_reference_encoder,
+            self.args.n_hidden_conformer_encoder,
+        )
+        self.p_norm = nn.InstanceNorm1d(
+            self.args.bottleneck_size_p_reference_encoder,
+        )
+        self.decoder = Conformer(
+            dim=self.args.n_hidden_conformer_decoder,
+            n_layers=self.args.n_layers_conformer_decoder,
+            n_heads=self.args.n_heads_conformer_decoder,
+            speaker_embedding_dim=self.embedded_speaker_dim,
+            p_dropout=self.args.dropout_conformer_decoder,
+            kernel_size_conv_mod=self.args.kernel_size_conv_mod_conformer_decoder,
+            lrelu_slope=self.args.lrelu_slope,
+        )
+
+        padding_idx = self.tokenizer.characters.pad_id
+        self.src_word_emb = EmbeddingPadded(
+            self.args.num_chars, self.args.n_hidden_conformer_encoder, padding_idx=padding_idx
+        )
+        self.to_mel = nn.Linear(
+            self.args.n_hidden_conformer_decoder,
+            self.args.num_mels,
+        )
+
+        self.energy_scaler = torch.nn.BatchNorm1d(1, affine=False, track_running_stats=True, momentum=None)
+        self.energy_scaler.requires_grad_(False)
+
+    def init_multispeaker(self, args: Coqpit):  # pylint: disable=unused-argument
+        """Init for multi-speaker training."""
+        self.embedded_speaker_dim = 0
+        self.num_speakers = self.args.num_speakers
+        self.audio_transform = None
+
+        if self.speaker_manager:
+            self.num_speakers = self.speaker_manager.num_speakers
+
+        if self.args.use_speaker_embedding:
+            self._init_speaker_embedding()
+
+        if self.args.use_d_vector_file:
+            self._init_d_vector()
+
+    @staticmethod
+    def _set_cond_input(aux_input: Dict):
+        """Set the speaker conditioning input based on the multi-speaker mode."""
+        sid, g, lid, durations = None, None, None, None
+        if "speaker_ids" in aux_input and aux_input["speaker_ids"] is not None:
+            sid = aux_input["speaker_ids"]
+            if sid.ndim == 0:
+                sid = sid.unsqueeze_(0)
+        if "d_vectors" in aux_input and aux_input["d_vectors"] is not None:
+            g = F.normalize(aux_input["d_vectors"])  # .unsqueeze_(-1)
+            if g.ndim == 2:
+                g = g  #  .unsqueeze_(0) # pylint: disable=self-assigning-variable
+
+        if "durations" in aux_input and aux_input["durations"] is not None:
+            durations = aux_input["durations"]
+
+        return sid, g, lid, durations
+
+    def get_aux_input(self, aux_input: Dict):
+        sid, g, lid, _ = self._set_cond_input(aux_input)
+        return {"speaker_ids": sid, "style_wav": None, "d_vectors": g, "language_ids": lid}
+
+    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 set_embedding_dims(self):
+    #     if self.embedded_speaker_dim > 0:
+    #         self.embedding_dims = self.embedded_speaker_dim
+    #     else:
+    #         self.embedding_dims = 0
+
+    def _init_speaker_embedding(self):
+        # pylint: disable=attribute-defined-outside-init
+        if self.num_speakers > 0:
+            print(" > initialization of speaker-embedding layers.")
+            self.embedded_speaker_dim = self.args.speaker_embedding_channels
+            self.emb_g = nn.Embedding(self.num_speakers, self.embedded_speaker_dim)
+
+    def _init_d_vector(self):
+        # pylint: disable=attribute-defined-outside-init
+        if hasattr(self, "emb_g"):
+            raise ValueError("[!] Speaker embedding layer already initialized before d_vector settings.")
+        self.embedded_speaker_dim = self.args.d_vector_dim
+
+    @staticmethod
+    def generate_attn(dr, x_mask, y_mask=None):
+        """Generate an attention mask from the linear scale durations.
+
+        Args:
+            dr (Tensor): Linear scale durations.
+            x_mask (Tensor): Mask for the input (character) sequence.
+            y_mask (Tensor): Mask for the output (spectrogram) sequence. Compute it from the predicted durations
+                if None. Defaults to None.
+
+        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_with_durations(
+        self,
+        o_en: torch.FloatTensor,
+        dr: torch.IntTensor,
+        x_mask: torch.IntTensor,
+        y_lengths: torch.IntTensor,
+    ):
+        y_mask = torch.unsqueeze(sequence_mask(y_lengths, None), 1).to(o_en.dtype)
+        attn = self.generate_attn(dr, x_mask, y_mask)
+        o_en_ex = torch.einsum("kmn, kjm -> kjn", [attn.float(), o_en])
+        return y_mask, o_en_ex, attn.transpose(1, 2)
+
+    def _forward_aligner(
+        self,
+        x: torch.FloatTensor,
+        y: torch.FloatTensor,
+        x_mask: torch.IntTensor,
+        y_mask: torch.IntTensor,
+        attn_priors: torch.FloatTensor,
+    ) -> 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.
+            attn_priors (torch.FloatTensor): Prior for the aligner network map.
+
+        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]`
+            - attn_priors: :math:`[B, T_de, T_en]`
+
+            - aligner_durations: :math:`[B, T_en]`
+            - aligner_soft: :math:`[B, T_de, T_en]`
+            - aligner_logprob: :math:`[B, 1, T_de, T_en]`
+            - aligner_mas: :math:`[B, T_de, T_en]`
+        """
+        attn_mask = torch.unsqueeze(x_mask, -1) * torch.unsqueeze(y_mask, 2)  # [B, 1, T_en, T_de]
+        aligner_soft, aligner_logprob = self.aligner(y.transpose(1, 2), x.transpose(1, 2), x_mask, attn_priors)
+        aligner_mas = maximum_path(
+            aligner_soft.squeeze(1).transpose(1, 2).contiguous(), attn_mask.squeeze(1).contiguous()
+        )
+        aligner_durations = torch.sum(aligner_mas, -1).int()
+        aligner_soft = aligner_soft.squeeze(1)  # [B, T_max2, T_max]
+        aligner_mas = aligner_mas.transpose(1, 2)  # [B, T_max, T_max2] -> [B, T_max2, T_max]
+        return aligner_durations, aligner_soft, aligner_logprob, aligner_mas
+
+    def average_utterance_prosody(  # pylint: disable=no-self-use
+        self, u_prosody_pred: torch.Tensor, src_mask: torch.Tensor
+    ) -> torch.Tensor:
+        lengths = ((~src_mask) * 1.0).sum(1)
+        u_prosody_pred = u_prosody_pred.sum(1, keepdim=True) / lengths.view(-1, 1, 1)
+        return u_prosody_pred
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        src_lens: torch.Tensor,
+        mels: torch.Tensor,
+        mel_lens: torch.Tensor,
+        pitches: torch.Tensor,
+        energies: torch.Tensor,
+        attn_priors: torch.Tensor,
+        use_ground_truth: bool = True,
+        d_vectors: torch.Tensor = None,
+        speaker_idx: torch.Tensor = None,
+    ) -> Dict[str, torch.Tensor]:
+        sid, g, lid, _ = self._set_cond_input(  # pylint: disable=unused-variable
+            {"d_vectors": d_vectors, "speaker_ids": speaker_idx}
+        )  # pylint: disable=unused-variable
+
+        src_mask = get_mask_from_lengths(src_lens)  # [B, T_src]
+        mel_mask = get_mask_from_lengths(mel_lens)  # [B, T_mel]
+
+        # Token embeddings
+        token_embeddings = self.src_word_emb(tokens)  # [B, T_src, C_hidden]
+        token_embeddings = token_embeddings.masked_fill(src_mask.unsqueeze(-1), 0.0)
+
+        # Alignment network and durations
+        aligner_durations, aligner_soft, aligner_logprob, aligner_mas = self._forward_aligner(
+            x=token_embeddings,
+            y=mels.transpose(1, 2),
+            x_mask=~src_mask[:, None],
+            y_mask=~mel_mask[:, None],
+            attn_priors=attn_priors,
+        )
+        dr = aligner_durations  # [B, T_en]
+
+        # Embeddings
+        speaker_embedding = None
+        if d_vectors is not None:
+            speaker_embedding = g
+        elif speaker_idx is not None:
+            speaker_embedding = F.normalize(self.emb_g(sid))
+
+        pos_encoding = positional_encoding(
+            self.emb_dim,
+            max(token_embeddings.shape[1], max(mel_lens)),
+            device=token_embeddings.device,
+        )
+        encoder_outputs = self.encoder(
+            token_embeddings,
+            src_mask,
+            speaker_embedding=speaker_embedding,
+            encoding=pos_encoding,
+        )
+
+        u_prosody_ref = self.u_norm(self.utterance_prosody_encoder(mels=mels, mel_lens=mel_lens))
+        u_prosody_pred = self.u_norm(
+            self.average_utterance_prosody(
+                u_prosody_pred=self.utterance_prosody_predictor(x=encoder_outputs, mask=src_mask),
+                src_mask=src_mask,
+            )
+        )
+
+        if use_ground_truth:
+            encoder_outputs = encoder_outputs + self.u_bottle_out(u_prosody_ref)
+        else:
+            encoder_outputs = encoder_outputs + self.u_bottle_out(u_prosody_pred)
+
+        p_prosody_ref = self.p_norm(
+            self.phoneme_prosody_encoder(
+                x=encoder_outputs, src_mask=src_mask, mels=mels, mel_lens=mel_lens, encoding=pos_encoding
+            )
+        )
+        p_prosody_pred = self.p_norm(self.phoneme_prosody_predictor(x=encoder_outputs, mask=src_mask))
+
+        if use_ground_truth:
+            encoder_outputs = encoder_outputs + self.p_bottle_out(p_prosody_ref)
+        else:
+            encoder_outputs = encoder_outputs + self.p_bottle_out(p_prosody_pred)
+
+        encoder_outputs_res = encoder_outputs
+
+        pitch_pred, avg_pitch_target, pitch_emb = self.pitch_adaptor.get_pitch_embedding_train(
+            x=encoder_outputs,
+            target=pitches,
+            dr=dr,
+            mask=src_mask,
+        )
+
+        energy_pred, avg_energy_target, energy_emb = self.energy_adaptor.get_energy_embedding_train(
+            x=encoder_outputs,
+            target=energies,
+            dr=dr,
+            mask=src_mask,
+        )
+
+        encoder_outputs = encoder_outputs.transpose(1, 2) + pitch_emb + energy_emb
+        log_duration_prediction = self.duration_predictor(x=encoder_outputs_res.detach(), mask=src_mask)
+
+        mel_pred_mask, encoder_outputs_ex, alignments = self._expand_encoder_with_durations(
+            o_en=encoder_outputs, y_lengths=mel_lens, dr=dr, x_mask=~src_mask[:, None]
+        )
+
+        x = self.decoder(
+            encoder_outputs_ex.transpose(1, 2),
+            mel_mask,
+            speaker_embedding=speaker_embedding,
+            encoding=pos_encoding,
+        )
+        x = self.to_mel(x)
+
+        dr = torch.log(dr + 1)
+
+        dr_pred = torch.exp(log_duration_prediction) - 1
+        alignments_dp = self.generate_attn(dr_pred, src_mask.unsqueeze(1), mel_pred_mask)  # [B, T_max, T_max2']
+
+        return {
+            "model_outputs": x,
+            "pitch_pred": pitch_pred,
+            "pitch_target": avg_pitch_target,
+            "energy_pred": energy_pred,
+            "energy_target": avg_energy_target,
+            "u_prosody_pred": u_prosody_pred,
+            "u_prosody_ref": u_prosody_ref,
+            "p_prosody_pred": p_prosody_pred,
+            "p_prosody_ref": p_prosody_ref,
+            "alignments_dp": alignments_dp,
+            "alignments": alignments,  # [B, T_de, T_en]
+            "aligner_soft": aligner_soft,
+            "aligner_mas": aligner_mas,
+            "aligner_durations": aligner_durations,
+            "aligner_logprob": aligner_logprob,
+            "dr_log_pred": log_duration_prediction.squeeze(1),  # [B, T]
+            "dr_log_target": dr.squeeze(1),  # [B, T]
+            "spk_emb": speaker_embedding,
+        }
+
+    @torch.no_grad()
+    def inference(
+        self,
+        tokens: torch.Tensor,
+        speaker_idx: torch.Tensor,
+        p_control: float = None,  # TODO # pylint: disable=unused-argument
+        d_control: float = None,  # TODO # pylint: disable=unused-argument
+        d_vectors: torch.Tensor = None,
+        pitch_transform: Callable = None,
+        energy_transform: Callable = None,
+    ) -> torch.Tensor:
+        src_mask = get_mask_from_lengths(torch.tensor([tokens.shape[1]], dtype=torch.int64, device=tokens.device))
+        src_lens = torch.tensor(tokens.shape[1:2]).to(tokens.device)  # pylint: disable=unused-variable
+        sid, g, lid, _ = self._set_cond_input(  # pylint: disable=unused-variable
+            {"d_vectors": d_vectors, "speaker_ids": speaker_idx}
+        )  # pylint: disable=unused-variable
+
+        token_embeddings = self.src_word_emb(tokens)
+        token_embeddings = token_embeddings.masked_fill(src_mask.unsqueeze(-1), 0.0)
+
+        # Embeddings
+        speaker_embedding = None
+        if d_vectors is not None:
+            speaker_embedding = g
+        elif speaker_idx is not None:
+            speaker_embedding = F.normalize(self.emb_g(sid))
+
+        pos_encoding = positional_encoding(
+            self.emb_dim,
+            token_embeddings.shape[1],
+            device=token_embeddings.device,
+        )
+        encoder_outputs = self.encoder(
+            token_embeddings,
+            src_mask,
+            speaker_embedding=speaker_embedding,
+            encoding=pos_encoding,
+        )
+
+        u_prosody_pred = self.u_norm(
+            self.average_utterance_prosody(
+                u_prosody_pred=self.utterance_prosody_predictor(x=encoder_outputs, mask=src_mask),
+                src_mask=src_mask,
+            )
+        )
+        encoder_outputs = encoder_outputs + self.u_bottle_out(u_prosody_pred).expand_as(encoder_outputs)
+
+        p_prosody_pred = self.p_norm(
+            self.phoneme_prosody_predictor(
+                x=encoder_outputs,
+                mask=src_mask,
+            )
+        )
+        encoder_outputs = encoder_outputs + self.p_bottle_out(p_prosody_pred).expand_as(encoder_outputs)
+
+        encoder_outputs_res = encoder_outputs
+
+        pitch_emb_pred, pitch_pred = self.pitch_adaptor.get_pitch_embedding(
+            x=encoder_outputs,
+            mask=src_mask,
+            pitch_transform=pitch_transform,
+            pitch_mean=self.pitch_mean if hasattr(self, "pitch_mean") else None,
+            pitch_std=self.pitch_std if hasattr(self, "pitch_std") else None,
+        )
+
+        energy_emb_pred, energy_pred = self.energy_adaptor.get_energy_embedding(
+            x=encoder_outputs, mask=src_mask, energy_transform=energy_transform
+        )
+        encoder_outputs = encoder_outputs.transpose(1, 2) + pitch_emb_pred + energy_emb_pred
+
+        log_duration_pred = self.duration_predictor(
+            x=encoder_outputs_res.detach(), mask=src_mask
+        )  # [B, C_hidden, T_src] -> [B, T_src]
+        duration_pred = (torch.exp(log_duration_pred) - 1) * (~src_mask) * self.length_scale  # -> [B, T_src]
+        duration_pred[duration_pred < 1] = 1.0  # -> [B, T_src]
+        duration_pred = torch.round(duration_pred)  # -> [B, T_src]
+        mel_lens = duration_pred.sum(1)  # -> [B,]
+
+        _, encoder_outputs_ex, alignments = self._expand_encoder_with_durations(
+            o_en=encoder_outputs, y_lengths=mel_lens, dr=duration_pred.squeeze(1), x_mask=~src_mask[:, None]
+        )
+
+        mel_mask = get_mask_from_lengths(
+            torch.tensor([encoder_outputs_ex.shape[2]], dtype=torch.int64, device=encoder_outputs_ex.device)
+        )
+
+        if encoder_outputs_ex.shape[1] > pos_encoding.shape[1]:
+            encoding = positional_encoding(self.emb_dim, encoder_outputs_ex.shape[2], device=tokens.device)
+
+        # [B, C_hidden, T_src], [B, 1, T_src], [B, C_emb], [B, T_src, C_hidden] -> [B, C_hidden, T_src]
+        x = self.decoder(
+            encoder_outputs_ex.transpose(1, 2),
+            mel_mask,
+            speaker_embedding=speaker_embedding,
+            encoding=encoding,
+        )
+        x = self.to_mel(x)
+        outputs = {
+            "model_outputs": x,
+            "alignments": alignments,
+            # "pitch": pitch_emb_pred,
+            "durations": duration_pred,
+            "pitch": pitch_pred,
+            "energy": energy_pred,
+            "spk_emb": speaker_embedding,
+        }
+        return outputs

TTS/tts/layers/delightful_tts/conformer.py

@@ -0,0 +1,450 @@
+### credit: https://github.com/dunky11/voicesmith
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn  # pylint: disable=consider-using-from-import
+import torch.nn.functional as F
+
+from TTS.tts.layers.delightful_tts.conv_layers import Conv1dGLU, DepthWiseConv1d, PointwiseConv1d
+from TTS.tts.layers.delightful_tts.networks import GLUActivation
+
+
+def calc_same_padding(kernel_size: int) -> Tuple[int, int]:
+    pad = kernel_size // 2
+    return (pad, pad - (kernel_size + 1) % 2)
+
+
+class Conformer(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        n_layers: int,
+        n_heads: int,
+        speaker_embedding_dim: int,
+        p_dropout: float,
+        kernel_size_conv_mod: int,
+        lrelu_slope: float,
+    ):
+        """
+        A Transformer variant that integrates both CNNs and Transformers components.
+        Conformer proposes a novel combination of self-attention and convolution, in which self-attention
+        learns the global interaction while the convolutions efficiently capture the local correlations.
+
+        Args:
+            dim (int): Number of the dimensions for the model.
+            n_layers (int): Number of model layers.
+            n_heads (int): The number of attention heads.
+            speaker_embedding_dim (int): Number of speaker embedding dimensions.
+            p_dropout (float): Probabilty of dropout.
+            kernel_size_conv_mod (int): Size of kernels for convolution modules.
+
+        Inputs: inputs, mask
+            - **inputs** (batch, time, dim): Tensor containing input vector
+            - **encoding** (batch, time, dim): Positional embedding tensor
+            - **mask** (batch, 1, time2) or (batch, time1, time2): Tensor containing indices to be masked
+        Returns:
+            - **outputs** (batch, time, dim): Tensor produced by Conformer Encoder.
+        """
+        super().__init__()
+        d_k = d_v = dim // n_heads
+        self.layer_stack = nn.ModuleList(
+            [
+                ConformerBlock(
+                    dim,
+                    n_heads,
+                    d_k,
+                    d_v,
+                    kernel_size_conv_mod=kernel_size_conv_mod,
+                    dropout=p_dropout,
+                    speaker_embedding_dim=speaker_embedding_dim,
+                    lrelu_slope=lrelu_slope,
+                )
+                for _ in range(n_layers)
+            ]
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor,
+        speaker_embedding: torch.Tensor,
+        encoding: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Shapes:
+            - x: :math:`[B, T_src, C]`
+            - mask: :math: `[B]`
+            - speaker_embedding: :math: `[B, C]`
+            - encoding: :math: `[B, T_max2, C]`
+        """
+
+        attn_mask = mask.view((mask.shape[0], 1, 1, mask.shape[1]))
+        for enc_layer in self.layer_stack:
+            x = enc_layer(
+                x,
+                mask=mask,
+                slf_attn_mask=attn_mask,
+                speaker_embedding=speaker_embedding,
+                encoding=encoding,
+            )
+        return x
+
+
+class ConformerBlock(torch.nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        n_head: int,
+        d_k: int,  # pylint: disable=unused-argument
+        d_v: int,  # pylint: disable=unused-argument
+        kernel_size_conv_mod: int,
+        speaker_embedding_dim: int,
+        dropout: float,
+        lrelu_slope: float = 0.3,
+    ):
+        """
+        A Conformer block is composed of four modules stacked together,
+        A feed-forward module, a self-attention module, a convolution module,
+        and a second feed-forward module in the end. The block starts with two Feed forward
+        modules sandwiching the Multi-Headed Self-Attention module and the Conv module.
+
+        Args:
+            d_model (int): The dimension of model
+            n_head (int): The number of attention heads.
+            kernel_size_conv_mod (int): Size of kernels for convolution modules.
+            speaker_embedding_dim (int): Number of speaker embedding dimensions.
+            emotion_embedding_dim (int): Number of emotion embedding dimensions.
+            dropout (float): Probabilty of dropout.
+
+        Inputs: inputs, mask
+            - **inputs** (batch, time, dim): Tensor containing input vector
+            - **encoding** (batch, time, dim): Positional embedding tensor
+            - **slf_attn_mask** (batch, 1, 1, time1): Tensor containing indices to be masked in self attention module
+            - **mask** (batch, 1, time2) or (batch, time1, time2): Tensor containing indices to be masked
+        Returns:
+            - **outputs** (batch, time, dim): Tensor produced by the Conformer Block.
+        """
+        super().__init__()
+        if isinstance(speaker_embedding_dim, int):
+            self.conditioning = Conv1dGLU(
+                d_model=d_model,
+                kernel_size=kernel_size_conv_mod,
+                padding=kernel_size_conv_mod // 2,
+                embedding_dim=speaker_embedding_dim,
+            )
+
+        self.ff = FeedForward(d_model=d_model, dropout=dropout, kernel_size=3, lrelu_slope=lrelu_slope)
+        self.conformer_conv_1 = ConformerConvModule(
+            d_model, kernel_size=kernel_size_conv_mod, dropout=dropout, lrelu_slope=lrelu_slope
+        )
+        self.ln = nn.LayerNorm(d_model)
+        self.slf_attn = ConformerMultiHeadedSelfAttention(d_model=d_model, num_heads=n_head, dropout_p=dropout)
+        self.conformer_conv_2 = ConformerConvModule(
+            d_model, kernel_size=kernel_size_conv_mod, dropout=dropout, lrelu_slope=lrelu_slope
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        speaker_embedding: torch.Tensor,
+        mask: torch.Tensor,
+        slf_attn_mask: torch.Tensor,
+        encoding: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Shapes:
+            - x: :math:`[B, T_src, C]`
+            - mask: :math: `[B]`
+            - slf_attn_mask: :math: `[B, 1, 1, T_src]`
+            - speaker_embedding: :math: `[B, C]`
+            - emotion_embedding: :math: `[B, C]`
+            - encoding: :math: `[B, T_max2, C]`
+        """
+        if speaker_embedding is not None:
+            x = self.conditioning(x, embeddings=speaker_embedding)
+        x = self.ff(x) + x
+        x = self.conformer_conv_1(x) + x
+        res = x
+        x = self.ln(x)
+        x, _ = self.slf_attn(query=x, key=x, value=x, mask=slf_attn_mask, encoding=encoding)
+        x = x + res
+        x = x.masked_fill(mask.unsqueeze(-1), 0)
+
+        x = self.conformer_conv_2(x) + x
+        return x
+
+
+class FeedForward(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        kernel_size: int,
+        dropout: float,
+        lrelu_slope: float,
+        expansion_factor: int = 4,
+    ):
+        """
+        Feed Forward module for conformer block.
+
+        Args:
+            d_model (int): The dimension of model.
+            kernel_size (int): Size of the kernels for conv layers.
+            dropout (float): probability of dropout.
+            expansion_factor (int): The factor by which to project the number of channels.
+            lrelu_slope (int): the negative slope factor for the leaky relu activation.
+
+        Inputs: inputs
+            - **inputs** (batch, time, dim): Tensor containing input vector
+        Returns:
+            - **outputs** (batch, time, dim): Tensor produced by the feed forward module.
+        """
+        super().__init__()
+        self.dropout = nn.Dropout(dropout)
+        self.ln = nn.LayerNorm(d_model)
+        self.conv_1 = nn.Conv1d(
+            d_model,
+            d_model * expansion_factor,
+            kernel_size=kernel_size,
+            padding=kernel_size // 2,
+        )
+        self.act = nn.LeakyReLU(lrelu_slope)
+        self.conv_2 = nn.Conv1d(d_model * expansion_factor, d_model, kernel_size=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Shapes:
+            x: :math: `[B, T, C]`
+        """
+        x = self.ln(x)
+        x = x.permute((0, 2, 1))
+        x = self.conv_1(x)
+        x = x.permute((0, 2, 1))
+        x = self.act(x)
+        x = self.dropout(x)
+        x = x.permute((0, 2, 1))
+        x = self.conv_2(x)
+        x = x.permute((0, 2, 1))
+        x = self.dropout(x)
+        x = 0.5 * x
+        return x
+
+
+class ConformerConvModule(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        expansion_factor: int = 2,
+        kernel_size: int = 7,
+        dropout: float = 0.1,
+        lrelu_slope: float = 0.3,
+    ):
+        """
+        Convolution module for conformer. Starts with a gating machanism.
+        a pointwise convolution and a gated linear unit (GLU). This is followed
+        by a single 1-D depthwise convolution layer. Batchnorm is deployed just after the convolution
+        to help with training. it also contains an expansion factor to project the number of channels.
+
+        Args:
+            d_model (int): The dimension of model.
+            expansion_factor (int): The factor by which to project the number of channels.
+            kernel_size (int): Size of kernels for convolution modules.
+            dropout (float): Probabilty of dropout.
+            lrelu_slope (float): The slope coefficient for leaky relu activation.
+
+        Inputs: inputs
+            - **inputs** (batch, time, dim): Tensor containing input vector
+        Returns:
+            - **outputs** (batch, time, dim): Tensor produced by the conv module.
+
+        """
+        super().__init__()
+        inner_dim = d_model * expansion_factor
+        self.ln_1 = nn.LayerNorm(d_model)
+        self.conv_1 = PointwiseConv1d(d_model, inner_dim * 2)
+        self.conv_act = GLUActivation(slope=lrelu_slope)
+        self.depthwise = DepthWiseConv1d(
+            inner_dim,
+            inner_dim,
+            kernel_size=kernel_size,
+            padding=calc_same_padding(kernel_size)[0],
+        )
+        self.ln_2 = nn.GroupNorm(1, inner_dim)
+        self.activation = nn.LeakyReLU(lrelu_slope)
+        self.conv_2 = PointwiseConv1d(inner_dim, d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Shapes:
+            x: :math: `[B, T, C]`
+        """
+        x = self.ln_1(x)
+        x = x.permute(0, 2, 1)
+        x = self.conv_1(x)
+        x = self.conv_act(x)
+        x = self.depthwise(x)
+        x = self.ln_2(x)
+        x = self.activation(x)
+        x = self.conv_2(x)
+        x = x.permute(0, 2, 1)
+        x = self.dropout(x)
+        return x
+
+
+class ConformerMultiHeadedSelfAttention(nn.Module):
+    """
+    Conformer employ multi-headed self-attention (MHSA) while integrating an important technique from Transformer-XL,
+    the relative sinusoidal positional encoding scheme. The relative positional encoding allows the self-attention
+    module to generalize better on different input length and the resulting encoder is more robust to the variance of
+    the utterance length. Conformer use prenorm residual units with dropout which helps training
+    and regularizing deeper models.
+    Args:
+        d_model (int): The dimension of model
+        num_heads (int): The number of attention heads.
+        dropout_p (float): probability of dropout
+    Inputs: inputs, mask
+        - **inputs** (batch, time, dim): Tensor containing input vector
+        - **mask** (batch, 1, time2) or (batch, time1, time2): Tensor containing indices to be masked
+    Returns:
+        - **outputs** (batch, time, dim): Tensor produces by relative multi headed self attention module.
+    """
+
+    def __init__(self, d_model: int, num_heads: int, dropout_p: float):
+        super().__init__()
+        self.attention = RelativeMultiHeadAttention(d_model=d_model, num_heads=num_heads)
+        self.dropout = nn.Dropout(p=dropout_p)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        mask: torch.Tensor,
+        encoding: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size, seq_length, _ = key.size()  # pylint: disable=unused-variable
+        encoding = encoding[:, : key.shape[1]]
+        encoding = encoding.repeat(batch_size, 1, 1)
+        outputs, attn = self.attention(query, key, value, pos_embedding=encoding, mask=mask)
+        outputs = self.dropout(outputs)
+        return outputs, attn
+
+
+class RelativeMultiHeadAttention(nn.Module):
+    """
+    Multi-head attention with relative positional encoding.
+    This concept was proposed in the "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context"
+    Args:
+        d_model (int): The dimension of model
+        num_heads (int): The number of attention heads.
+    Inputs: query, key, value, pos_embedding, mask
+        - **query** (batch, time, dim): Tensor containing query vector
+        - **key** (batch, time, dim): Tensor containing key vector
+        - **value** (batch, time, dim): Tensor containing value vector
+        - **pos_embedding** (batch, time, dim): Positional embedding tensor
+        - **mask** (batch, 1, time2) or (batch, time1, time2): Tensor containing indices to be masked
+    Returns:
+        - **outputs**: Tensor produces by relative multi head attention module.
+    """
+
+    def __init__(
+        self,
+        d_model: int = 512,
+        num_heads: int = 16,
+    ):
+        super().__init__()
+        assert d_model % num_heads == 0, "d_model % num_heads should be zero."
+        self.d_model = d_model
+        self.d_head = int(d_model / num_heads)
+        self.num_heads = num_heads
+        self.sqrt_dim = math.sqrt(d_model)
+
+        self.query_proj = nn.Linear(d_model, d_model)
+        self.key_proj = nn.Linear(d_model, d_model, bias=False)
+        self.value_proj = nn.Linear(d_model, d_model, bias=False)
+        self.pos_proj = nn.Linear(d_model, d_model, bias=False)
+
+        self.u_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))
+        self.v_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))
+        torch.nn.init.xavier_uniform_(self.u_bias)
+        torch.nn.init.xavier_uniform_(self.v_bias)
+        self.out_proj = nn.Linear(d_model, d_model)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        pos_embedding: torch.Tensor,
+        mask: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size = query.shape[0]
+        query = self.query_proj(query).view(batch_size, -1, self.num_heads, self.d_head)
+        key = self.key_proj(key).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)
+        value = self.value_proj(value).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)
+        pos_embedding = self.pos_proj(pos_embedding).view(batch_size, -1, self.num_heads, self.d_head)
+        u_bias = self.u_bias.expand_as(query)
+        v_bias = self.v_bias.expand_as(query)
+        a = (query + u_bias).transpose(1, 2)
+        content_score = a @ key.transpose(2, 3)
+        b = (query + v_bias).transpose(1, 2)
+        pos_score = b @ pos_embedding.permute(0, 2, 3, 1)
+        pos_score = self._relative_shift(pos_score)
+
+        score = content_score + pos_score
+        score = score * (1.0 / self.sqrt_dim)
+
+        score.masked_fill_(mask, -1e9)
+
+        attn = F.softmax(score, -1)
+
+        context = (attn @ value).transpose(1, 2)
+        context = context.contiguous().view(batch_size, -1, self.d_model)
+
+        return self.out_proj(context), attn
+
+    def _relative_shift(self, pos_score: torch.Tensor) -> torch.Tensor:  # pylint: disable=no-self-use
+        batch_size, num_heads, seq_length1, seq_length2 = pos_score.size()
+        zeros = torch.zeros((batch_size, num_heads, seq_length1, 1), device=pos_score.device)
+        padded_pos_score = torch.cat([zeros, pos_score], dim=-1)
+        padded_pos_score = padded_pos_score.view(batch_size, num_heads, seq_length2 + 1, seq_length1)
+        pos_score = padded_pos_score[:, :, 1:].view_as(pos_score)
+        return pos_score
+
+
+class MultiHeadAttention(nn.Module):
+    """
+    input:
+        query --- [N, T_q, query_dim]
+        key --- [N, T_k, key_dim]
+    output:
+        out --- [N, T_q, num_units]
+    """
+
+    def __init__(self, query_dim: int, key_dim: int, num_units: int, num_heads: int):
+        super().__init__()
+        self.num_units = num_units
+        self.num_heads = num_heads
+        self.key_dim = key_dim
+
+        self.W_query = nn.Linear(in_features=query_dim, out_features=num_units, bias=False)
+        self.W_key = nn.Linear(in_features=key_dim, out_features=num_units, bias=False)
+        self.W_value = nn.Linear(in_features=key_dim, out_features=num_units, bias=False)
+
+    def forward(self, query: torch.Tensor, key: torch.Tensor) -> torch.Tensor:
+        querys = self.W_query(query)  # [N, T_q, num_units]
+        keys = self.W_key(key)  # [N, T_k, num_units]
+        values = self.W_value(key)
+        split_size = self.num_units // self.num_heads
+        querys = torch.stack(torch.split(querys, split_size, dim=2), dim=0)  # [h, N, T_q, num_units/h]
+        keys = torch.stack(torch.split(keys, split_size, dim=2), dim=0)  # [h, N, T_k, num_units/h]
+        values = torch.stack(torch.split(values, split_size, dim=2), dim=0)  # [h, N, T_k, num_units/h]
+        # score = softmax(QK^T / (d_k ** 0.5))
+        scores = torch.matmul(querys, keys.transpose(2, 3))  # [h, N, T_q, T_k]
+        scores = scores / (self.key_dim**0.5)
+        scores = F.softmax(scores, dim=3)
+        # out = score * V
+        out = torch.matmul(scores, values)  # [h, N, T_q, num_units/h]
+        out = torch.cat(torch.split(out, 1, dim=0), dim=3).squeeze(0)  # [N, T_q, num_units]
+        return out

TTS/tts/layers/delightful_tts/conv_layers.py

@@ -0,0 +1,671 @@
+from typing import Tuple
+
+import torch
+import torch.nn as nn  # pylint: disable=consider-using-from-import
+import torch.nn.functional as F
+from torch.nn.utils import parametrize
+
+from TTS.tts.layers.delightful_tts.kernel_predictor import KernelPredictor
+
+
+def calc_same_padding(kernel_size: int) -> Tuple[int, int]:
+    pad = kernel_size // 2
+    return (pad, pad - (kernel_size + 1) % 2)
+
+
+class ConvNorm(nn.Module):
+    """A 1-dimensional convolutional layer with optional weight normalization.
+
+    This layer wraps a 1D convolutional layer from PyTorch and applies
+    optional weight normalization. The layer can be used in a similar way to
+    the convolutional layers in PyTorch's `torch.nn` module.
+
+    Args:
+        in_channels (int): The number of channels in the input signal.
+        out_channels (int): The number of channels in the output signal.
+        kernel_size (int, optional): The size of the convolving kernel.
+            Defaults to 1.
+        stride (int, optional): The stride of the convolution. Defaults to 1.
+        padding (int, optional): Zero-padding added to both sides of the input.
+            If `None`, the padding will be calculated so that the output has
+            the same length as the input. Defaults to `None`.
+        dilation (int, optional): Spacing between kernel elements. Defaults to 1.
+        bias (bool, optional): If `True`, add bias after convolution. Defaults to `True`.
+        w_init_gain (str, optional): The weight initialization function to use.
+            Can be either 'linear' or 'relu'. Defaults to 'linear'.
+        use_weight_norm (bool, optional): If `True`, apply weight normalization
+            to the convolutional weights. Defaults to `False`.
+
+    Shapes:
+     - Input: :math:`[N, D, T]`
+
+    - Output: :math:`[N, out_dim, T]` where `out_dim` is the number of output dimensions.
+
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=1,
+        stride=1,
+        padding=None,
+        dilation=1,
+        bias=True,
+        w_init_gain="linear",
+        use_weight_norm=False,
+    ):
+        super(ConvNorm, self).__init__()  # pylint: disable=super-with-arguments
+        if padding is None:
+            assert kernel_size % 2 == 1
+            padding = int(dilation * (kernel_size - 1) / 2)
+        self.kernel_size = kernel_size
+        self.dilation = dilation
+        self.use_weight_norm = use_weight_norm
+        conv_fn = nn.Conv1d
+        self.conv = conv_fn(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+        )
+        nn.init.xavier_uniform_(self.conv.weight, gain=nn.init.calculate_gain(w_init_gain))
+        if self.use_weight_norm:
+            self.conv = nn.utils.parametrizations.weight_norm(self.conv)
+
+    def forward(self, signal, mask=None):
+        conv_signal = self.conv(signal)
+        if mask is not None:
+            # always re-zero output if mask is
+            # available to match zero-padding
+            conv_signal = conv_signal * mask
+        return conv_signal
+
+
+class ConvLSTMLinear(nn.Module):
+    def __init__(
+        self,
+        in_dim,
+        out_dim,
+        n_layers=2,
+        n_channels=256,
+        kernel_size=3,
+        p_dropout=0.1,
+        lstm_type="bilstm",
+        use_linear=True,
+    ):
+        super(ConvLSTMLinear, self).__init__()  # pylint: disable=super-with-arguments
+        self.out_dim = out_dim
+        self.lstm_type = lstm_type
+        self.use_linear = use_linear
+        self.dropout = nn.Dropout(p=p_dropout)
+
+        convolutions = []
+        for i in range(n_layers):
+            conv_layer = ConvNorm(
+                in_dim if i == 0 else n_channels,
+                n_channels,
+                kernel_size=kernel_size,
+                stride=1,
+                padding=int((kernel_size - 1) / 2),
+                dilation=1,
+                w_init_gain="relu",
+            )
+            conv_layer = nn.utils.parametrizations.weight_norm(conv_layer.conv, name="weight")
+            convolutions.append(conv_layer)
+
+        self.convolutions = nn.ModuleList(convolutions)
+
+        if not self.use_linear:
+            n_channels = out_dim
+
+        if self.lstm_type != "":
+            use_bilstm = False
+            lstm_channels = n_channels
+            if self.lstm_type == "bilstm":
+                use_bilstm = True
+                lstm_channels = int(n_channels // 2)
+
+            self.bilstm = nn.LSTM(n_channels, lstm_channels, 1, batch_first=True, bidirectional=use_bilstm)
+            lstm_norm_fn_pntr = nn.utils.spectral_norm
+            self.bilstm = lstm_norm_fn_pntr(self.bilstm, "weight_hh_l0")
+            if self.lstm_type == "bilstm":
+                self.bilstm = lstm_norm_fn_pntr(self.bilstm, "weight_hh_l0_reverse")
+
+        if self.use_linear:
+            self.dense = nn.Linear(n_channels, out_dim)
+
+    def run_padded_sequence(self, context, lens):
+        context_embedded = []
+        for b_ind in range(context.size()[0]):  # TODO: speed up
+            curr_context = context[b_ind : b_ind + 1, :, : lens[b_ind]].clone()
+            for conv in self.convolutions:
+                curr_context = self.dropout(F.relu(conv(curr_context)))
+            context_embedded.append(curr_context[0].transpose(0, 1))
+        context = nn.utils.rnn.pad_sequence(context_embedded, batch_first=True)
+        return context
+
+    def run_unsorted_inputs(self, fn, context, lens):  # pylint: disable=no-self-use
+        lens_sorted, ids_sorted = torch.sort(lens, descending=True)
+        unsort_ids = [0] * lens.size(0)
+        for i in range(len(ids_sorted)):  # pylint: disable=consider-using-enumerate
+            unsort_ids[ids_sorted[i]] = i
+        lens_sorted = lens_sorted.long().cpu()
+
+        context = context[ids_sorted]
+        context = nn.utils.rnn.pack_padded_sequence(context, lens_sorted, batch_first=True)
+        context = fn(context)[0]
+        context = nn.utils.rnn.pad_packed_sequence(context, batch_first=True)[0]
+
+        # map back to original indices
+        context = context[unsort_ids]
+        return context
+
+    def forward(self, context, lens):
+        if context.size()[0] > 1:
+            context = self.run_padded_sequence(context, lens)
+            # to B, D, T
+            context = context.transpose(1, 2)
+        else:
+            for conv in self.convolutions:
+                context = self.dropout(F.relu(conv(context)))
+
+        if self.lstm_type != "":
+            context = context.transpose(1, 2)
+            self.bilstm.flatten_parameters()
+            if lens is not None:
+                context = self.run_unsorted_inputs(self.bilstm, context, lens)
+            else:
+                context = self.bilstm(context)[0]
+            context = context.transpose(1, 2)
+
+        x_hat = context
+        if self.use_linear:
+            x_hat = self.dense(context.transpose(1, 2)).transpose(1, 2)
+
+        return x_hat
+
+
+class DepthWiseConv1d(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, padding: int):
+        super().__init__()
+        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, padding=padding, groups=in_channels)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)
+
+
+class PointwiseConv1d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        stride: int = 1,
+        padding: int = 0,
+        bias: bool = True,
+    ):
+        super().__init__()
+        self.conv = nn.Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=stride,
+            padding=padding,
+            bias=bias,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)
+
+
+class BSConv1d(nn.Module):
+    """https://arxiv.org/pdf/2003.13549.pdf"""
+
+    def __init__(self, channels_in: int, channels_out: int, kernel_size: int, padding: int):
+        super().__init__()
+        self.pointwise = nn.Conv1d(channels_in, channels_out, kernel_size=1)
+        self.depthwise = nn.Conv1d(
+            channels_out,
+            channels_out,
+            kernel_size=kernel_size,
+            padding=padding,
+            groups=channels_out,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x1 = self.pointwise(x)
+        x2 = self.depthwise(x1)
+        return x2
+
+
+class BSConv2d(nn.Module):
+    """https://arxiv.org/pdf/2003.13549.pdf"""
+
+    def __init__(self, channels_in: int, channels_out: int, kernel_size: int, padding: int):
+        super().__init__()
+        self.pointwise = nn.Conv2d(channels_in, channels_out, kernel_size=1)
+        self.depthwise = nn.Conv2d(
+            channels_out,
+            channels_out,
+            kernel_size=kernel_size,
+            padding=padding,
+            groups=channels_out,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x1 = self.pointwise(x)
+        x2 = self.depthwise(x1)
+        return x2
+
+
+class Conv1dGLU(nn.Module):
+    """From DeepVoice 3"""
+
+    def __init__(self, d_model: int, kernel_size: int, padding: int, embedding_dim: int):
+        super().__init__()
+        self.conv = BSConv1d(d_model, 2 * d_model, kernel_size=kernel_size, padding=padding)
+        self.embedding_proj = nn.Linear(embedding_dim, d_model)
+        self.register_buffer("sqrt", torch.sqrt(torch.FloatTensor([0.5])).squeeze(0))
+        self.softsign = torch.nn.Softsign()
+
+    def forward(self, x: torch.Tensor, embeddings: torch.Tensor) -> torch.Tensor:
+        x = x.permute((0, 2, 1))
+        residual = x
+        x = self.conv(x)
+        splitdim = 1
+        a, b = x.split(x.size(splitdim) // 2, dim=splitdim)
+        embeddings = self.embedding_proj(embeddings).unsqueeze(2)
+        softsign = self.softsign(embeddings)
+        softsign = softsign.expand_as(a)
+        a = a + softsign
+        x = a * torch.sigmoid(b)
+        x = x + residual
+        x = x * self.sqrt
+        x = x.permute((0, 2, 1))
+        return x
+
+
+class ConvTransposed(nn.Module):
+    """
+    A 1D convolutional transposed layer for PyTorch.
+    This layer applies a 1D convolutional transpose operation to its input tensor,
+    where the number of channels of the input tensor is the same as the number of channels of the output tensor.
+
+    Attributes:
+        in_channels (int): The number of channels in the input tensor.
+        out_channels (int): The number of channels in the output tensor.
+        kernel_size (int): The size of the convolutional kernel. Default: 1.
+        padding (int): The number of padding elements to add to the input tensor. Default: 0.
+        conv (BSConv1d): The 1D convolutional transpose layer.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 1,
+        padding: int = 0,
+    ):
+        super().__init__()
+        self.conv = BSConv1d(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            padding=padding,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.contiguous().transpose(1, 2)
+        x = self.conv(x)
+        x = x.contiguous().transpose(1, 2)
+        return x
+
+
+class DepthwiseConvModule(nn.Module):
+    def __init__(self, dim: int, kernel_size: int = 7, expansion: int = 4, lrelu_slope: float = 0.3):
+        super().__init__()
+        padding = calc_same_padding(kernel_size)
+        self.depthwise = nn.Conv1d(
+            dim,
+            dim * expansion,
+            kernel_size=kernel_size,
+            padding=padding[0],
+            groups=dim,
+        )
+        self.act = nn.LeakyReLU(lrelu_slope)
+        self.out = nn.Conv1d(dim * expansion, dim, 1, 1, 0)
+        self.ln = nn.LayerNorm(dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.ln(x)
+        x = x.permute((0, 2, 1))
+        x = self.depthwise(x)
+        x = self.act(x)
+        x = self.out(x)
+        x = x.permute((0, 2, 1))
+        return x
+
+
+class AddCoords(nn.Module):
+    def __init__(self, rank: int, with_r: bool = False):
+        super().__init__()
+        self.rank = rank
+        self.with_r = with_r
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.rank == 1:
+            batch_size_shape, channel_in_shape, dim_x = x.shape  # pylint: disable=unused-variable
+            xx_range = torch.arange(dim_x, dtype=torch.int32)
+            xx_channel = xx_range[None, None, :]
+
+            xx_channel = xx_channel.float() / (dim_x - 1)
+            xx_channel = xx_channel * 2 - 1
+            xx_channel = xx_channel.repeat(batch_size_shape, 1, 1)
+
+            xx_channel = xx_channel.to(x.device)
+            out = torch.cat([x, xx_channel], dim=1)
+
+            if self.with_r:
+                rr = torch.sqrt(torch.pow(xx_channel - 0.5, 2))
+                out = torch.cat([out, rr], dim=1)
+
+        elif self.rank == 2:
+            batch_size_shape, channel_in_shape, dim_y, dim_x = x.shape
+            xx_ones = torch.ones([1, 1, 1, dim_x], dtype=torch.int32)
+            yy_ones = torch.ones([1, 1, 1, dim_y], dtype=torch.int32)
+
+            xx_range = torch.arange(dim_y, dtype=torch.int32)
+            yy_range = torch.arange(dim_x, dtype=torch.int32)
+            xx_range = xx_range[None, None, :, None]
+            yy_range = yy_range[None, None, :, None]
+
+            xx_channel = torch.matmul(xx_range, xx_ones)
+            yy_channel = torch.matmul(yy_range, yy_ones)
+
+            # transpose y
+            yy_channel = yy_channel.permute(0, 1, 3, 2)
+
+            xx_channel = xx_channel.float() / (dim_y - 1)
+            yy_channel = yy_channel.float() / (dim_x - 1)
+
+            xx_channel = xx_channel * 2 - 1
+            yy_channel = yy_channel * 2 - 1
+
+            xx_channel = xx_channel.repeat(batch_size_shape, 1, 1, 1)
+            yy_channel = yy_channel.repeat(batch_size_shape, 1, 1, 1)
+
+            xx_channel = xx_channel.to(x.device)
+            yy_channel = yy_channel.to(x.device)
+
+            out = torch.cat([x, xx_channel, yy_channel], dim=1)
+
+            if self.with_r:
+                rr = torch.sqrt(torch.pow(xx_channel - 0.5, 2) + torch.pow(yy_channel - 0.5, 2))
+                out = torch.cat([out, rr], dim=1)
+
+        elif self.rank == 3:
+            batch_size_shape, channel_in_shape, dim_z, dim_y, dim_x = x.shape
+            xx_ones = torch.ones([1, 1, 1, 1, dim_x], dtype=torch.int32)
+            yy_ones = torch.ones([1, 1, 1, 1, dim_y], dtype=torch.int32)
+            zz_ones = torch.ones([1, 1, 1, 1, dim_z], dtype=torch.int32)
+
+            xy_range = torch.arange(dim_y, dtype=torch.int32)
+            xy_range = xy_range[None, None, None, :, None]
+
+            yz_range = torch.arange(dim_z, dtype=torch.int32)
+            yz_range = yz_range[None, None, None, :, None]
+
+            zx_range = torch.arange(dim_x, dtype=torch.int32)
+            zx_range = zx_range[None, None, None, :, None]
+
+            xy_channel = torch.matmul(xy_range, xx_ones)
+            xx_channel = torch.cat([xy_channel + i for i in range(dim_z)], dim=2)
+
+            yz_channel = torch.matmul(yz_range, yy_ones)
+            yz_channel = yz_channel.permute(0, 1, 3, 4, 2)
+            yy_channel = torch.cat([yz_channel + i for i in range(dim_x)], dim=4)
+
+            zx_channel = torch.matmul(zx_range, zz_ones)
+            zx_channel = zx_channel.permute(0, 1, 4, 2, 3)
+            zz_channel = torch.cat([zx_channel + i for i in range(dim_y)], dim=3)
+
+            xx_channel = xx_channel.to(x.device)
+            yy_channel = yy_channel.to(x.device)
+            zz_channel = zz_channel.to(x.device)
+            out = torch.cat([x, xx_channel, yy_channel, zz_channel], dim=1)
+
+            if self.with_r:
+                rr = torch.sqrt(
+                    torch.pow(xx_channel - 0.5, 2) + torch.pow(yy_channel - 0.5, 2) + torch.pow(zz_channel - 0.5, 2)
+                )
+                out = torch.cat([out, rr], dim=1)
+        else:
+            raise NotImplementedError
+
+        return out
+
+
+class CoordConv1d(nn.modules.conv.Conv1d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        padding: int = 0,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+        with_r: bool = False,
+    ):
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            groups,
+            bias,
+        )
+        self.rank = 1
+        self.addcoords = AddCoords(self.rank, with_r)
+        self.conv = nn.Conv1d(
+            in_channels + self.rank + int(with_r),
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            groups,
+            bias,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.addcoords(x)
+        x = self.conv(x)
+        return x
+
+
+class CoordConv2d(nn.modules.conv.Conv2d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        padding: int = 0,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+        with_r: bool = False,
+    ):
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            groups,
+            bias,
+        )
+        self.rank = 2
+        self.addcoords = AddCoords(self.rank, with_r)
+        self.conv = nn.Conv2d(
+            in_channels + self.rank + int(with_r),
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            groups,
+            bias,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.addcoords(x)
+        x = self.conv(x)
+        return x
+
+
+class LVCBlock(torch.nn.Module):
+    """the location-variable convolutions"""
+
+    def __init__(  # pylint: disable=dangerous-default-value
+        self,
+        in_channels,
+        cond_channels,
+        stride,
+        dilations=[1, 3, 9, 27],
+        lReLU_slope=0.2,
+        conv_kernel_size=3,
+        cond_hop_length=256,
+        kpnet_hidden_channels=64,
+        kpnet_conv_size=3,
+        kpnet_dropout=0.0,
+    ):
+        super().__init__()
+
+        self.cond_hop_length = cond_hop_length
+        self.conv_layers = len(dilations)
+        self.conv_kernel_size = conv_kernel_size
+
+        self.kernel_predictor = KernelPredictor(
+            cond_channels=cond_channels,
+            conv_in_channels=in_channels,
+            conv_out_channels=2 * in_channels,
+            conv_layers=len(dilations),
+            conv_kernel_size=conv_kernel_size,
+            kpnet_hidden_channels=kpnet_hidden_channels,
+            kpnet_conv_size=kpnet_conv_size,
+            kpnet_dropout=kpnet_dropout,
+            kpnet_nonlinear_activation_params={"negative_slope": lReLU_slope},
+        )
+
+        self.convt_pre = nn.Sequential(
+            nn.LeakyReLU(lReLU_slope),
+            nn.utils.parametrizations.weight_norm(
+                nn.ConvTranspose1d(
+                    in_channels,
+                    in_channels,
+                    2 * stride,
+                    stride=stride,
+                    padding=stride // 2 + stride % 2,
+                    output_padding=stride % 2,
+                )
+            ),
+        )
+
+        self.conv_blocks = nn.ModuleList()
+        for dilation in dilations:
+            self.conv_blocks.append(
+                nn.Sequential(
+                    nn.LeakyReLU(lReLU_slope),
+                    nn.utils.parametrizations.weight_norm(
+                        nn.Conv1d(
+                            in_channels,
+                            in_channels,
+                            conv_kernel_size,
+                            padding=dilation * (conv_kernel_size - 1) // 2,
+                            dilation=dilation,
+                        )
+                    ),
+                    nn.LeakyReLU(lReLU_slope),
+                )
+            )
+
+    def forward(self, x, c):
+        """forward propagation of the location-variable convolutions.
+        Args:
+            x (Tensor): the input sequence (batch, in_channels, in_length)
+            c (Tensor): the conditioning sequence (batch, cond_channels, cond_length)
+
+        Returns:
+            Tensor: the output sequence (batch, in_channels, in_length)
+        """
+        _, in_channels, _ = x.shape  # (B, c_g, L')
+
+        x = self.convt_pre(x)  # (B, c_g, stride * L')
+        kernels, bias = self.kernel_predictor(c)
+
+        for i, conv in enumerate(self.conv_blocks):
+            output = conv(x)  # (B, c_g, stride * L')
+
+            k = kernels[:, i, :, :, :, :]  # (B, 2 * c_g, c_g, kernel_size, cond_length)
+            b = bias[:, i, :, :]  # (B, 2 * c_g, cond_length)
+
+            output = self.location_variable_convolution(
+                output, k, b, hop_size=self.cond_hop_length
+            )  # (B, 2 * c_g, stride * L'): LVC
+            x = x + torch.sigmoid(output[:, :in_channels, :]) * torch.tanh(
+                output[:, in_channels:, :]
+            )  # (B, c_g, stride * L'): GAU
+
+        return x
+
+    def location_variable_convolution(self, x, kernel, bias, dilation=1, hop_size=256):  # pylint: disable=no-self-use
+        """perform location-variable convolution operation on the input sequence (x) using the local convolution kernl.
+        Time: 414 μs ± 309 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each), test on NVIDIA V100.
+        Args:
+            x (Tensor): the input sequence (batch, in_channels, in_length).
+            kernel (Tensor): the local convolution kernel (batch, in_channel, out_channels, kernel_size, kernel_length)
+            bias (Tensor): the bias for the local convolution (batch, out_channels, kernel_length)
+            dilation (int): the dilation of convolution.
+            hop_size (int): the hop_size of the conditioning sequence.
+        Returns:
+            (Tensor): the output sequence after performing local convolution. (batch, out_channels, in_length).
+        """
+        batch, _, in_length = x.shape
+        batch, _, out_channels, kernel_size, kernel_length = kernel.shape
+        assert in_length == (kernel_length * hop_size), "length of (x, kernel) is not matched"
+
+        padding = dilation * int((kernel_size - 1) / 2)
+        x = F.pad(x, (padding, padding), "constant", 0)  # (batch, in_channels, in_length + 2*padding)
+        x = x.unfold(2, hop_size + 2 * padding, hop_size)  # (batch, in_channels, kernel_length, hop_size + 2*padding)
+
+        if hop_size < dilation:
+            x = F.pad(x, (0, dilation), "constant", 0)
+        x = x.unfold(
+            3, dilation, dilation
+        )  # (batch, in_channels, kernel_length, (hop_size + 2*padding)/dilation, dilation)
+        x = x[:, :, :, :, :hop_size]
+        x = x.transpose(3, 4)  # (batch, in_channels, kernel_length, dilation, (hop_size + 2*padding)/dilation)
+        x = x.unfold(4, kernel_size, 1)  # (batch, in_channels, kernel_length, dilation, _, kernel_size)
+
+        o = torch.einsum("bildsk,biokl->bolsd", x, kernel)
+        o = o.to(memory_format=torch.channels_last_3d)
+        bias = bias.unsqueeze(-1).unsqueeze(-1).to(memory_format=torch.channels_last_3d)
+        o = o + bias
+        o = o.contiguous().view(batch, out_channels, -1)
+
+        return o
+
+    def remove_weight_norm(self):
+        self.kernel_predictor.remove_weight_norm()
+        parametrize.remove_parametrizations(self.convt_pre[1], "weight")
+        for block in self.conv_blocks:
+            parametrize.remove_parametrizations(block[1], "weight")

TTS/tts/layers/delightful_tts/encoders.py

@@ -0,0 +1,261 @@
+from typing import List, Tuple, Union
+
+import torch
+import torch.nn as nn  # pylint: disable=consider-using-from-import
+import torch.nn.functional as F
+
+from TTS.tts.layers.delightful_tts.conformer import ConformerMultiHeadedSelfAttention
+from TTS.tts.layers.delightful_tts.conv_layers import CoordConv1d
+from TTS.tts.layers.delightful_tts.networks import STL
+
+
+def get_mask_from_lengths(lengths: torch.Tensor) -> torch.Tensor:
+    batch_size = lengths.shape[0]
+    max_len = torch.max(lengths).item()
+    ids = torch.arange(0, max_len, device=lengths.device).unsqueeze(0).expand(batch_size, -1)
+    mask = ids >= lengths.unsqueeze(1).expand(-1, max_len)
+    return mask
+
+
+def stride_lens(lens: torch.Tensor, stride: int = 2) -> torch.Tensor:
+    return torch.ceil(lens / stride).int()
+
+
+class ReferenceEncoder(nn.Module):
+    """
+    Referance encoder for utterance and phoneme prosody encoders. Reference encoder
+    made up of convolution and RNN layers.
+
+    Args:
+        num_mels (int): Number of mel frames to produce.
+        ref_enc_filters (list[int]): List of channel sizes for encoder layers.
+        ref_enc_size (int): Size of the kernel for the conv layers.
+        ref_enc_strides (List[int]): List of strides to use for conv layers.
+        ref_enc_gru_size (int): Number of hidden features for the gated recurrent unit.
+
+    Inputs: inputs, mask
+        - **inputs** (batch, dim, time): Tensor containing mel vector
+        - **lengths** (batch): Tensor containing the mel lengths.
+    Returns:
+        - **outputs** (batch, time, dim): Tensor produced by Reference Encoder.
+    """
+
+    def __init__(
+        self,
+        num_mels: int,
+        ref_enc_filters: List[Union[int, int, int, int, int, int]],
+        ref_enc_size: int,
+        ref_enc_strides: List[Union[int, int, int, int, int]],
+        ref_enc_gru_size: int,
+    ):
+        super().__init__()
+
+        n_mel_channels = num_mels
+        self.n_mel_channels = n_mel_channels
+        K = len(ref_enc_filters)
+        filters = [self.n_mel_channels] + ref_enc_filters
+        strides = [1] + ref_enc_strides
+        # Use CoordConv at the first layer to better preserve positional information: https://arxiv.org/pdf/1811.02122.pdf
+        convs = [
+            CoordConv1d(
+                in_channels=filters[0],
+                out_channels=filters[0 + 1],
+                kernel_size=ref_enc_size,
+                stride=strides[0],
+                padding=ref_enc_size // 2,
+                with_r=True,
+            )
+        ]
+        convs2 = [
+            nn.Conv1d(
+                in_channels=filters[i],
+                out_channels=filters[i + 1],
+                kernel_size=ref_enc_size,
+                stride=strides[i],
+                padding=ref_enc_size // 2,
+            )
+            for i in range(1, K)
+        ]
+        convs.extend(convs2)
+        self.convs = nn.ModuleList(convs)
+
+        self.norms = nn.ModuleList([nn.InstanceNorm1d(num_features=ref_enc_filters[i], affine=True) for i in range(K)])
+
+        self.gru = nn.GRU(
+            input_size=ref_enc_filters[-1],
+            hidden_size=ref_enc_gru_size,
+            batch_first=True,
+        )
+
+    def forward(self, x: torch.Tensor, mel_lens: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        inputs --- [N,  n_mels, timesteps]
+        outputs --- [N, E//2]
+        """
+
+        mel_masks = get_mask_from_lengths(mel_lens).unsqueeze(1)
+        x = x.masked_fill(mel_masks, 0)
+        for conv, norm in zip(self.convs, self.norms):
+            x = conv(x)
+            x = F.leaky_relu(x, 0.3)  # [N, 128, Ty//2^K, n_mels//2^K]
+            x = norm(x)
+
+        for _ in range(2):
+            mel_lens = stride_lens(mel_lens)
+
+        mel_masks = get_mask_from_lengths(mel_lens)
+
+        x = x.masked_fill(mel_masks.unsqueeze(1), 0)
+        x = x.permute((0, 2, 1))
+        x = torch.nn.utils.rnn.pack_padded_sequence(x, mel_lens.cpu().int(), batch_first=True, enforce_sorted=False)
+
+        self.gru.flatten_parameters()
+        x, memory = self.gru(x)  # memory --- [N, Ty, E//2], out --- [1, N, E//2]
+        x, _ = torch.nn.utils.rnn.pad_packed_sequence(x, batch_first=True)
+
+        return x, memory, mel_masks
+
+    def calculate_channels(  # pylint: disable=no-self-use
+        self, L: int, kernel_size: int, stride: int, pad: int, n_convs: int
+    ) -> int:
+        for _ in range(n_convs):
+            L = (L - kernel_size + 2 * pad) // stride + 1
+        return L
+
+
+class UtteranceLevelProsodyEncoder(nn.Module):
+    def __init__(
+        self,
+        num_mels: int,
+        ref_enc_filters: List[Union[int, int, int, int, int, int]],
+        ref_enc_size: int,
+        ref_enc_strides: List[Union[int, int, int, int, int]],
+        ref_enc_gru_size: int,
+        dropout: float,
+        n_hidden: int,
+        bottleneck_size_u: int,
+        token_num: int,
+    ):
+        """
+        Encoder to extract prosody from utterance. it is made up of a reference encoder
+        with a couple of linear layers and style token layer with dropout.
+
+        Args:
+            num_mels (int): Number of mel frames to produce.
+            ref_enc_filters (list[int]): List of channel sizes for ref encoder layers.
+            ref_enc_size (int): Size of the kernel for the ref encoder conv layers.
+            ref_enc_strides (List[int]): List of strides to use for teh ref encoder conv layers.
+            ref_enc_gru_size (int): Number of hidden features for the gated recurrent unit.
+            dropout (float): Probability of dropout.
+            n_hidden (int): Size of hidden layers.
+            bottleneck_size_u (int): Size of the bottle neck layer.
+
+        Inputs: inputs, mask
+            - **inputs** (batch, dim, time): Tensor containing mel vector
+            - **lengths** (batch): Tensor containing the mel lengths.
+        Returns:
+            - **outputs** (batch, 1, dim): Tensor produced by Utterance Level Prosody Encoder.
+        """
+        super().__init__()
+
+        self.E = n_hidden
+        self.d_q = self.d_k = n_hidden
+        bottleneck_size = bottleneck_size_u
+
+        self.encoder = ReferenceEncoder(
+            ref_enc_filters=ref_enc_filters,
+            ref_enc_gru_size=ref_enc_gru_size,
+            ref_enc_size=ref_enc_size,
+            ref_enc_strides=ref_enc_strides,
+            num_mels=num_mels,
+        )
+        self.encoder_prj = nn.Linear(ref_enc_gru_size, self.E // 2)
+        self.stl = STL(n_hidden=n_hidden, token_num=token_num)
+        self.encoder_bottleneck = nn.Linear(self.E, bottleneck_size)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, mels: torch.Tensor, mel_lens: torch.Tensor) -> torch.Tensor:
+        """
+        Shapes:
+            mels: :math: `[B, C, T]`
+            mel_lens: :math: `[B]`
+
+        out --- [N, seq_len, E]
+        """
+        _, embedded_prosody, _ = self.encoder(mels, mel_lens)
+
+        # Bottleneck
+        embedded_prosody = self.encoder_prj(embedded_prosody)
+
+        # Style Token
+        out = self.encoder_bottleneck(self.stl(embedded_prosody))
+        out = self.dropout(out)
+
+        out = out.view((-1, 1, out.shape[3]))
+        return out
+
+
+class PhonemeLevelProsodyEncoder(nn.Module):
+    def __init__(
+        self,
+        num_mels: int,
+        ref_enc_filters: List[Union[int, int, int, int, int, int]],
+        ref_enc_size: int,
+        ref_enc_strides: List[Union[int, int, int, int, int]],
+        ref_enc_gru_size: int,
+        dropout: float,
+        n_hidden: int,
+        n_heads: int,
+        bottleneck_size_p: int,
+    ):
+        super().__init__()
+
+        self.E = n_hidden
+        self.d_q = self.d_k = n_hidden
+        bottleneck_size = bottleneck_size_p
+
+        self.encoder = ReferenceEncoder(
+            ref_enc_filters=ref_enc_filters,
+            ref_enc_gru_size=ref_enc_gru_size,
+            ref_enc_size=ref_enc_size,
+            ref_enc_strides=ref_enc_strides,
+            num_mels=num_mels,
+        )
+        self.encoder_prj = nn.Linear(ref_enc_gru_size, n_hidden)
+        self.attention = ConformerMultiHeadedSelfAttention(
+            d_model=n_hidden,
+            num_heads=n_heads,
+            dropout_p=dropout,
+        )
+        self.encoder_bottleneck = nn.Linear(n_hidden, bottleneck_size)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        src_mask: torch.Tensor,
+        mels: torch.Tensor,
+        mel_lens: torch.Tensor,
+        encoding: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        x --- [N, seq_len, encoder_embedding_dim]
+        mels --- [N, Ty/r, n_mels*r], r=1
+        out --- [N, seq_len, bottleneck_size]
+        attn --- [N, seq_len, ref_len], Ty/r = ref_len
+        """
+        embedded_prosody, _, mel_masks = self.encoder(mels, mel_lens)
+
+        # Bottleneck
+        embedded_prosody = self.encoder_prj(embedded_prosody)
+
+        attn_mask = mel_masks.view((mel_masks.shape[0], 1, 1, -1))
+        x, _ = self.attention(
+            query=x,
+            key=embedded_prosody,
+            value=embedded_prosody,
+            mask=attn_mask,
+            encoding=encoding,
+        )
+        x = self.encoder_bottleneck(x)
+        x = x.masked_fill(src_mask.unsqueeze(-1), 0.0)
+        return x

TTS/tts/layers/delightful_tts/energy_adaptor.py

@@ -0,0 +1,82 @@
+from typing import Callable, Tuple
+
+import torch
+import torch.nn as nn  # pylint: disable=consider-using-from-import
+
+from TTS.tts.layers.delightful_tts.variance_predictor import VariancePredictor
+from TTS.tts.utils.helpers import average_over_durations
+
+
+class EnergyAdaptor(nn.Module):  # pylint: disable=abstract-method
+    """Variance Adaptor with an added 1D conv layer. Used to
+    get energy embeddings.
+
+    Args:
+        channels_in (int): Number of in channels for conv layers.
+        channels_out (int): Number of out channels.
+        kernel_size (int): Size the kernel for the conv layers.
+        dropout (float): Probability of dropout.
+        lrelu_slope (float): Slope for the leaky relu.
+        emb_kernel_size (int): Size the kernel for the pitch embedding.
+
+    Inputs: inputs, mask
+        - **inputs** (batch, time1, dim): Tensor containing input vector
+        - **target** (batch, 1, time2): Tensor containing the energy target
+        - **dr** (batch, time1): Tensor containing aligner durations vector
+        - **mask** (batch, time1): Tensor containing indices to be masked
+    Returns:
+        - **energy prediction** (batch, 1, time1): Tensor produced by energy predictor
+        - **energy embedding** (batch, channels, time1): Tensor produced energy adaptor
+        - **average energy target(train only)** (batch, 1, time1): Tensor produced after averaging over durations
+
+    """
+
+    def __init__(
+        self,
+        channels_in: int,
+        channels_hidden: int,
+        channels_out: int,
+        kernel_size: int,
+        dropout: float,
+        lrelu_slope: float,
+        emb_kernel_size: int,
+    ):
+        super().__init__()
+        self.energy_predictor = VariancePredictor(
+            channels_in=channels_in,
+            channels=channels_hidden,
+            channels_out=channels_out,
+            kernel_size=kernel_size,
+            p_dropout=dropout,
+            lrelu_slope=lrelu_slope,
+        )
+        self.energy_emb = nn.Conv1d(
+            1,
+            channels_hidden,
+            kernel_size=emb_kernel_size,
+            padding=int((emb_kernel_size - 1) / 2),
+        )
+
+    def get_energy_embedding_train(
+        self, x: torch.Tensor, target: torch.Tensor, dr: torch.IntTensor, mask: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Shapes:
+            x: :math: `[B, T_src, C]`
+            target: :math: `[B, 1, T_max2]`
+            dr: :math: `[B, T_src]`
+            mask: :math: `[B, T_src]`
+        """
+        energy_pred = self.energy_predictor(x, mask)
+        energy_pred.unsqueeze_(1)
+        avg_energy_target = average_over_durations(target, dr)
+        energy_emb = self.energy_emb(avg_energy_target)
+        return energy_pred, avg_energy_target, energy_emb
+
+    def get_energy_embedding(self, x: torch.Tensor, mask: torch.Tensor, energy_transform: Callable) -> torch.Tensor:
+        energy_pred = self.energy_predictor(x, mask)
+        energy_pred.unsqueeze_(1)
+        if energy_transform is not None:
+            energy_pred = energy_transform(energy_pred, (~mask).sum(dim=(1, 2)), self.pitch_mean, self.pitch_std)
+        energy_emb_pred = self.energy_emb(energy_pred)
+        return energy_emb_pred, energy_pred

TTS/tts/layers/delightful_tts/kernel_predictor.py

@@ -0,0 +1,128 @@
+import torch.nn as nn  # pylint: disable=consider-using-from-import
+from torch.nn.utils import parametrize
+
+
+class KernelPredictor(nn.Module):
+    """Kernel predictor for the location-variable convolutions
+
+    Args:
+            cond_channels (int): number of channel for the conditioning sequence,
+            conv_in_channels (int): number of channel for the input sequence,
+            conv_out_channels (int): number of channel for the output sequence,
+            conv_layers (int): number of layers
+
+    """
+
+    def __init__(  # pylint: disable=dangerous-default-value
+        self,
+        cond_channels,
+        conv_in_channels,
+        conv_out_channels,
+        conv_layers,
+        conv_kernel_size=3,
+        kpnet_hidden_channels=64,
+        kpnet_conv_size=3,
+        kpnet_dropout=0.0,
+        kpnet_nonlinear_activation="LeakyReLU",
+        kpnet_nonlinear_activation_params={"negative_slope": 0.1},
+    ):
+        super().__init__()
+
+        self.conv_in_channels = conv_in_channels
+        self.conv_out_channels = conv_out_channels
+        self.conv_kernel_size = conv_kernel_size
+        self.conv_layers = conv_layers
+
+        kpnet_kernel_channels = conv_in_channels * conv_out_channels * conv_kernel_size * conv_layers  # l_w
+        kpnet_bias_channels = conv_out_channels * conv_layers  # l_b
+
+        self.input_conv = nn.Sequential(
+            nn.utils.parametrizations.weight_norm(
+                nn.Conv1d(cond_channels, kpnet_hidden_channels, 5, padding=2, bias=True)
+            ),
+            getattr(nn, kpnet_nonlinear_activation)(**kpnet_nonlinear_activation_params),
+        )
+
+        self.residual_convs = nn.ModuleList()
+        padding = (kpnet_conv_size - 1) // 2
+        for _ in range(3):
+            self.residual_convs.append(
+                nn.Sequential(
+                    nn.Dropout(kpnet_dropout),
+                    nn.utils.parametrizations.weight_norm(
+                        nn.Conv1d(
+                            kpnet_hidden_channels,
+                            kpnet_hidden_channels,
+                            kpnet_conv_size,
+                            padding=padding,
+                            bias=True,
+                        )
+                    ),
+                    getattr(nn, kpnet_nonlinear_activation)(**kpnet_nonlinear_activation_params),
+                    nn.utils.parametrizations.weight_norm(
+                        nn.Conv1d(
+                            kpnet_hidden_channels,
+                            kpnet_hidden_channels,
+                            kpnet_conv_size,
+                            padding=padding,
+                            bias=True,
+                        )
+                    ),
+                    getattr(nn, kpnet_nonlinear_activation)(**kpnet_nonlinear_activation_params),
+                )
+            )
+        self.kernel_conv = nn.utils.parametrizations.weight_norm(
+            nn.Conv1d(
+                kpnet_hidden_channels,
+                kpnet_kernel_channels,
+                kpnet_conv_size,
+                padding=padding,
+                bias=True,
+            )
+        )
+        self.bias_conv = nn.utils.parametrizations.weight_norm(
+            nn.Conv1d(
+                kpnet_hidden_channels,
+                kpnet_bias_channels,
+                kpnet_conv_size,
+                padding=padding,
+                bias=True,
+            )
+        )
+
+    def forward(self, c):
+        """
+        Args:
+            c (Tensor): the conditioning sequence (batch, cond_channels, cond_length)
+        """
+        batch, _, cond_length = c.shape
+        c = self.input_conv(c)
+        for residual_conv in self.residual_convs:
+            residual_conv.to(c.device)
+            c = c + residual_conv(c)
+        k = self.kernel_conv(c)
+        b = self.bias_conv(c)
+        kernels = k.contiguous().view(
+            batch,
+            self.conv_layers,
+            self.conv_in_channels,
+            self.conv_out_channels,
+            self.conv_kernel_size,
+            cond_length,
+        )
+        bias = b.contiguous().view(
+            batch,
+            self.conv_layers,
+            self.conv_out_channels,
+            cond_length,
+        )
+
+        return kernels, bias
+
+    def remove_weight_norm(self):
+        parametrize.remove_parametrizations(self.input_conv[0], "weight")
+        parametrize.remove_parametrizations(self.kernel_conv, "weight")
+        parametrize.remove_parametrizations(self.bias_conv, "weight")
+        for block in self.residual_convs:
+            parametrize.remove_parametrizations(block[1], "weight")
+            parametrize.remove_parametrizations(block[3], "weight")

TTS/tts/layers/delightful_tts/networks.py

@@ -0,0 +1,219 @@
+import math
+from typing import Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn  # pylint: disable=consider-using-from-import
+import torch.nn.functional as F
+
+from TTS.tts.layers.delightful_tts.conv_layers import ConvNorm
+
+
+def initialize_embeddings(shape: Tuple[int]) -> torch.Tensor:
+    assert len(shape) == 2, "Can only initialize 2-D embedding matrices ..."
+    # Kaiming initialization
+    return torch.randn(shape) * np.sqrt(2 / shape[1])
+
+
+def positional_encoding(d_model: int, length: int, device: torch.device) -> torch.Tensor:
+    pe = torch.zeros(length, d_model, device=device)
+    position = torch.arange(0, length, dtype=torch.float, device=device).unsqueeze(1)
+    div_term = torch.exp(torch.arange(0, d_model, 2, device=device).float() * -(math.log(10000.0) / d_model))
+    pe[:, 0::2] = torch.sin(position * div_term)
+    pe[:, 1::2] = torch.cos(position * div_term)
+    pe = pe.unsqueeze(0)
+    return pe
+
+
+class BottleneckLayer(nn.Module):
+    """
+    Bottleneck layer for reducing the dimensionality of a tensor.
+
+    Args:
+        in_dim: The number of input dimensions.
+        reduction_factor: The factor by which to reduce the number of dimensions.
+        norm: The normalization method to use. Can be "weightnorm" or "instancenorm".
+        non_linearity: The non-linearity to use. Can be "relu" or "leakyrelu".
+        kernel_size: The size of the convolutional kernel.
+        use_partial_padding: Whether to use partial padding with the convolutional kernel.
+
+    Shape:
+        - Input: :math:`[N, in_dim]` where `N` is the batch size and `in_dim` is the number of input dimensions.
+
+        - Output: :math:`[N, out_dim]` where `out_dim` is the number of output dimensions.
+    """
+
+    def __init__(
+        self,
+        in_dim,
+        reduction_factor,
+        norm="weightnorm",
+        non_linearity="relu",
+        kernel_size=3,
+        use_partial_padding=False,  # pylint: disable=unused-argument
+    ):
+        super(BottleneckLayer, self).__init__()  # pylint: disable=super-with-arguments
+
+        self.reduction_factor = reduction_factor
+        reduced_dim = int(in_dim / reduction_factor)
+        self.out_dim = reduced_dim
+        if self.reduction_factor > 1:
+            fn = ConvNorm(in_dim, reduced_dim, kernel_size=kernel_size, use_weight_norm=(norm == "weightnorm"))
+            if norm == "instancenorm":
+                fn = nn.Sequential(fn, nn.InstanceNorm1d(reduced_dim, affine=True))
+
+            self.projection_fn = fn
+            self.non_linearity = nn.ReLU()
+            if non_linearity == "leakyrelu":
+                self.non_linearity = nn.LeakyReLU()
+
+    def forward(self, x):
+        if self.reduction_factor > 1:
+            x = self.projection_fn(x)
+            x = self.non_linearity(x)
+        return x
+
+
+class GLUActivation(nn.Module):
+    """Class that implements the Gated Linear Unit (GLU) activation function.
+
+    The GLU activation function is a variant of the Leaky ReLU activation function,
+    where the output of the activation function is gated by an input tensor.
+
+    """
+
+    def __init__(self, slope: float):
+        super().__init__()
+        self.lrelu = nn.LeakyReLU(slope)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out, gate = x.chunk(2, dim=1)
+        x = out * self.lrelu(gate)
+        return x
+
+
+class StyleEmbedAttention(nn.Module):
+    def __init__(self, query_dim: int, key_dim: int, num_units: int, num_heads: int):
+        super().__init__()
+        self.num_units = num_units
+        self.num_heads = num_heads
+        self.key_dim = key_dim
+
+        self.W_query = nn.Linear(in_features=query_dim, out_features=num_units, bias=False)
+        self.W_key = nn.Linear(in_features=key_dim, out_features=num_units, bias=False)
+        self.W_value = nn.Linear(in_features=key_dim, out_features=num_units, bias=False)
+
+    def forward(self, query: torch.Tensor, key_soft: torch.Tensor) -> torch.Tensor:
+        values = self.W_value(key_soft)
+        split_size = self.num_units // self.num_heads
+        values = torch.stack(torch.split(values, split_size, dim=2), dim=0)
+
+        out_soft = scores_soft = None
+        querys = self.W_query(query)  # [N, T_q, num_units]
+        keys = self.W_key(key_soft)  # [N, T_k, num_units]
+
+        # [h, N, T_q, num_units/h]
+        querys = torch.stack(torch.split(querys, split_size, dim=2), dim=0)
+        # [h, N, T_k, num_units/h]
+        keys = torch.stack(torch.split(keys, split_size, dim=2), dim=0)
+        # [h, N, T_k, num_units/h]
+
+        # score = softmax(QK^T / (d_k ** 0.5))
+        scores_soft = torch.matmul(querys, keys.transpose(2, 3))  # [h, N, T_q, T_k]
+        scores_soft = scores_soft / (self.key_dim**0.5)
+        scores_soft = F.softmax(scores_soft, dim=3)
+
+        # out = score * V
+        # [h, N, T_q, num_units/h]
+        out_soft = torch.matmul(scores_soft, values)
+        out_soft = torch.cat(torch.split(out_soft, 1, dim=0), dim=3).squeeze(0)  # [N, T_q, num_units]
+
+        return out_soft  # , scores_soft
+
+
+class EmbeddingPadded(nn.Module):
+    def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int):
+        super().__init__()
+        padding_mult = torch.ones((num_embeddings, 1), dtype=torch.int64)
+        padding_mult[padding_idx] = 0
+        self.register_buffer("padding_mult", padding_mult)
+        self.embeddings = nn.parameter.Parameter(initialize_embeddings((num_embeddings, embedding_dim)))
+
+    def forward(self, idx: torch.Tensor) -> torch.Tensor:
+        embeddings_zeroed = self.embeddings * self.padding_mult
+        x = F.embedding(idx, embeddings_zeroed)
+        return x
+
+
+class EmbeddingProjBlock(nn.Module):
+    def __init__(self, embedding_dim: int):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [
+                nn.Linear(embedding_dim, embedding_dim),
+                nn.LeakyReLU(0.3),
+                nn.Linear(embedding_dim, embedding_dim),
+                nn.LeakyReLU(0.3),
+            ]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        res = x
+        for layer in self.layers:
+            x = layer(x)
+        x = x + res
+        return x
+
+
+class LinearNorm(nn.Module):
+    def __init__(self, in_features: int, out_features: int, bias: bool = False):
+        super().__init__()
+        self.linear = nn.Linear(in_features, out_features, bias)
+
+        nn.init.xavier_uniform_(self.linear.weight)
+        if bias:
+            nn.init.constant_(self.linear.bias, 0.0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.linear(x)
+        return x
+
+
+class STL(nn.Module):
+    """
+    A PyTorch module for the Style Token Layer (STL) as described in
+    "A Style-Based Generator Architecture for Generative Adversarial Networks"
+    (https://arxiv.org/abs/1812.04948)
+
+    The STL applies a multi-headed attention mechanism over the learned style tokens,
+    using the text input as the query and the style tokens as the keys and values.
+    The output of the attention mechanism is used as the text's style embedding.
+
+    Args:
+        token_num (int): The number of style tokens.
+        n_hidden (int): Number of hidden dimensions.
+    """
+
+    def __init__(self, n_hidden: int, token_num: int):
+        super(STL, self).__init__()  # pylint: disable=super-with-arguments
+
+        num_heads = 1
+        E = n_hidden
+        self.token_num = token_num
+        self.embed = nn.Parameter(torch.FloatTensor(self.token_num, E // num_heads))
+        d_q = E // 2
+        d_k = E // num_heads
+        self.attention = StyleEmbedAttention(query_dim=d_q, key_dim=d_k, num_units=E, num_heads=num_heads)
+
+        torch.nn.init.normal_(self.embed, mean=0, std=0.5)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        N = x.size(0)
+        query = x.unsqueeze(1)  # [N, 1, E//2]
+
+        keys_soft = torch.tanh(self.embed).unsqueeze(0).expand(N, -1, -1)  # [N, token_num, E // num_heads]
+
+        # Weighted sum
+        emotion_embed_soft = self.attention(query, keys_soft)
+
+        return emotion_embed_soft

TTS/tts/layers/delightful_tts/phoneme_prosody_predictor.py

@@ -0,0 +1,65 @@
+import torch
+import torch.nn as nn  # pylint: disable=consider-using-from-import
+
+from TTS.tts.layers.delightful_tts.conv_layers import ConvTransposed
+
+
+class PhonemeProsodyPredictor(nn.Module):
+    """Non-parallel Prosody Predictor inspired by: https://arxiv.org/pdf/2102.00851.pdf
+    It consists of 2 layers of  1D convolutions each followed by a relu activation, layer norm
+    and dropout, then finally a linear layer.
+
+    Args:
+        hidden_size (int): Size of hidden channels.
+        kernel_size (int): Kernel size for the conv layers.
+        dropout: (float): Probability of dropout.
+        bottleneck_size (int): bottleneck size for last linear layer.
+        lrelu_slope (float): Slope of the leaky relu.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        kernel_size: int,
+        dropout: float,
+        bottleneck_size: int,
+        lrelu_slope: float,
+    ):
+        super().__init__()
+        self.d_model = hidden_size
+        self.layers = nn.ModuleList(
+            [
+                ConvTransposed(
+                    self.d_model,
+                    self.d_model,
+                    kernel_size=kernel_size,
+                    padding=(kernel_size - 1) // 2,
+                ),
+                nn.LeakyReLU(lrelu_slope),
+                nn.LayerNorm(self.d_model),
+                nn.Dropout(dropout),
+                ConvTransposed(
+                    self.d_model,
+                    self.d_model,
+                    kernel_size=kernel_size,
+                    padding=(kernel_size - 1) // 2,
+                ),
+                nn.LeakyReLU(lrelu_slope),
+                nn.LayerNorm(self.d_model),
+                nn.Dropout(dropout),
+            ]
+        )
+        self.predictor_bottleneck = nn.Linear(self.d_model, bottleneck_size)
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        """
+        Shapes:
+            x: :math: `[B, T, D]`
+            mask: :math: `[B, T]`
+        """
+        mask = mask.unsqueeze(2)
+        for layer in self.layers:
+            x = layer(x)
+        x = x.masked_fill(mask, 0.0)
+        x = self.predictor_bottleneck(x)
+        return x

TTS/tts/layers/delightful_tts/pitch_adaptor.py

@@ -0,0 +1,88 @@
+from typing import Callable, Tuple
+
+import torch
+import torch.nn as nn  # pylint: disable=consider-using-from-import
+
+from TTS.tts.layers.delightful_tts.variance_predictor import VariancePredictor
+from TTS.tts.utils.helpers import average_over_durations
+
+
+class PitchAdaptor(nn.Module):  # pylint: disable=abstract-method
+    """Module to get pitch embeddings via pitch predictor
+
+    Args:
+        n_input (int): Number of pitch predictor input channels.
+        n_hidden (int): Number of pitch predictor hidden channels.
+        n_out (int): Number of pitch predictor out channels.
+        kernel size (int): Size of the kernel for conv layers.
+        emb_kernel_size (int): Size the kernel for the pitch embedding.
+        p_dropout (float): Probability of dropout.
+        lrelu_slope (float): Slope for the leaky relu.
+
+    Inputs: inputs, mask
+        - **inputs** (batch, time1, dim): Tensor containing input vector
+        - **target** (batch, 1, time2): Tensor containing the pitch target
+        - **dr** (batch, time1): Tensor containing aligner durations vector
+        - **mask** (batch, time1): Tensor containing indices to be masked
+    Returns:
+        - **pitch prediction** (batch, 1, time1): Tensor produced by pitch predictor
+        - **pitch embedding** (batch, channels, time1): Tensor produced pitch pitch adaptor
+        - **average pitch target(train only)** (batch, 1, time1): Tensor produced after averaging over durations
+    """
+
+    def __init__(
+        self,
+        n_input: int,
+        n_hidden: int,
+        n_out: int,
+        kernel_size: int,
+        emb_kernel_size: int,
+        p_dropout: float,
+        lrelu_slope: float,
+    ):
+        super().__init__()
+        self.pitch_predictor = VariancePredictor(
+            channels_in=n_input,
+            channels=n_hidden,
+            channels_out=n_out,
+            kernel_size=kernel_size,
+            p_dropout=p_dropout,
+            lrelu_slope=lrelu_slope,
+        )
+        self.pitch_emb = nn.Conv1d(
+            1,
+            n_input,
+            kernel_size=emb_kernel_size,
+            padding=int((emb_kernel_size - 1) / 2),
+        )
+
+    def get_pitch_embedding_train(
+        self, x: torch.Tensor, target: torch.Tensor, dr: torch.IntTensor, mask: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Shapes:
+            x: :math: `[B, T_src, C]`
+            target: :math: `[B, 1, T_max2]`
+            dr: :math: `[B, T_src]`
+            mask: :math: `[B, T_src]`
+        """
+        pitch_pred = self.pitch_predictor(x, mask)  # [B, T_src, C_hidden], [B, T_src] --> [B, T_src]
+        pitch_pred.unsqueeze_(1)  # --> [B, 1, T_src]
+        avg_pitch_target = average_over_durations(target, dr)  # [B, 1, T_mel], [B, T_src] --> [B, 1, T_src]
+        pitch_emb = self.pitch_emb(avg_pitch_target)  # [B, 1, T_src] --> [B, C_hidden, T_src]
+        return pitch_pred, avg_pitch_target, pitch_emb
+
+    def get_pitch_embedding(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor,
+        pitch_transform: Callable,
+        pitch_mean: torch.Tensor,
+        pitch_std: torch.Tensor,
+    ) -> torch.Tensor:
+        pitch_pred = self.pitch_predictor(x, mask)
+        if pitch_transform is not None:
+            pitch_pred = pitch_transform(pitch_pred, (~mask).sum(), pitch_mean, pitch_std)
+        pitch_pred.unsqueeze_(1)
+        pitch_emb_pred = self.pitch_emb(pitch_pred)
+        return pitch_emb_pred, pitch_pred

TTS/tts/layers/delightful_tts/variance_predictor.py

@@ -0,0 +1,68 @@
+import torch
+import torch.nn as nn  # pylint: disable=consider-using-from-import
+
+from TTS.tts.layers.delightful_tts.conv_layers import ConvTransposed
+
+
+class VariancePredictor(nn.Module):
+    """
+    Network is 2-layer 1D convolutions with leaky relu activation and then
+    followed by layer normalization then a dropout layer and finally an
+    extra linear layer to project the hidden states into the output sequence.
+
+    Args:
+        channels_in (int): Number of in channels for conv layers.
+        channels_out (int): Number of out channels for the last linear layer.
+        kernel_size (int): Size the kernel for the conv layers.
+        p_dropout (float): Probability of dropout.
+        lrelu_slope (float): Slope for the leaky relu.
+
+    Inputs: inputs, mask
+        - **inputs** (batch, time, dim): Tensor containing input vector
+        - **mask** (batch, time): Tensor containing indices to be masked
+    Returns:
+        - **outputs** (batch, time): Tensor produced by last linear layer.
+    """
+
+    def __init__(
+        self, channels_in: int, channels: int, channels_out: int, kernel_size: int, p_dropout: float, lrelu_slope: float
+    ):
+        super().__init__()
+
+        self.layers = nn.ModuleList(
+            [
+                ConvTransposed(
+                    channels_in,
+                    channels,
+                    kernel_size=kernel_size,
+                    padding=(kernel_size - 1) // 2,
+                ),
+                nn.LeakyReLU(lrelu_slope),
+                nn.LayerNorm(channels),
+                nn.Dropout(p_dropout),
+                ConvTransposed(
+                    channels,
+                    channels,
+                    kernel_size=kernel_size,
+                    padding=(kernel_size - 1) // 2,
+                ),
+                nn.LeakyReLU(lrelu_slope),
+                nn.LayerNorm(channels),
+                nn.Dropout(p_dropout),
+            ]
+        )
+
+        self.linear_layer = nn.Linear(channels, channels_out)
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        """
+        Shapes:
+            x: :math: `[B, T_src, C]`
+            mask: :math: `[B, T_src]`
+        """
+        for layer in self.layers:
+            x = layer(x)
+        x = self.linear_layer(x)
+        x = x.squeeze(-1)
+        x = x.masked_fill(mask, 0.0)
+        return x

TTS/tts/layers/feed_forward/decoder.py

@@ -0,0 +1,228 @@
+import torch
+from torch import nn
+
+from TTS.tts.layers.generic.res_conv_bn import Conv1dBN, Conv1dBNBlock, ResidualConv1dBNBlock
+from TTS.tts.layers.generic.transformer import FFTransformerBlock
+from TTS.tts.layers.generic.wavenet import WNBlocks
+from TTS.tts.layers.glow_tts.transformer import RelativePositionTransformer
+
+
+class WaveNetDecoder(nn.Module):
+    """WaveNet based decoder with a prenet and a postnet.
+
+    prenet: conv1d_1x1
+    postnet: 3 x [conv1d_1x1 -> relu] -> conv1d_1x1
+
+    TODO: Integrate speaker conditioning vector.
+
+    Note:
+        default wavenet parameters;
+            params = {
+                "num_blocks": 12,
+                "hidden_channels":192,
+                "kernel_size": 5,
+                "dilation_rate": 1,
+                "num_layers": 4,
+                "dropout_p": 0.05
+            }
+
+    Args:
+        in_channels (int): number of input channels.
+        out_channels (int): number of output channels.
+        hidden_channels (int): number of hidden channels for prenet and postnet.
+        params (dict): dictionary for residual convolutional blocks.
+    """
+
+    def __init__(self, in_channels, out_channels, hidden_channels, c_in_channels, params):
+        super().__init__()
+        # prenet
+        self.prenet = torch.nn.Conv1d(in_channels, params["hidden_channels"], 1)
+        # wavenet layers
+        self.wn = WNBlocks(params["hidden_channels"], c_in_channels=c_in_channels, **params)
+        # postnet
+        self.postnet = [
+            torch.nn.Conv1d(params["hidden_channels"], hidden_channels, 1),
+            torch.nn.ReLU(),
+            torch.nn.Conv1d(hidden_channels, hidden_channels, 1),
+            torch.nn.ReLU(),
+            torch.nn.Conv1d(hidden_channels, hidden_channels, 1),
+            torch.nn.ReLU(),
+            torch.nn.Conv1d(hidden_channels, out_channels, 1),
+        ]
+        self.postnet = nn.Sequential(*self.postnet)
+
+    def forward(self, x, x_mask=None, g=None):
+        x = self.prenet(x) * x_mask
+        x = self.wn(x, x_mask, g)
+        o = self.postnet(x) * x_mask
+        return o
+
+
+class RelativePositionTransformerDecoder(nn.Module):
+    """Decoder with Relative Positional Transformer.
+
+    Note:
+        Default params
+            params={
+                'hidden_channels_ffn': 128,
+                'num_heads': 2,
+                "kernel_size": 3,
+                "dropout_p": 0.1,
+                "num_layers": 8,
+                "rel_attn_window_size": 4,
+                "input_length": None
+            }
+
+    Args:
+        in_channels (int): number of input channels.
+        out_channels (int): number of output channels.
+        hidden_channels (int): number of hidden channels including Transformer layers.
+        params (dict): dictionary for residual convolutional blocks.
+    """
+
+    def __init__(self, in_channels, out_channels, hidden_channels, params):
+        super().__init__()
+        self.prenet = Conv1dBN(in_channels, hidden_channels, 1, 1)
+        self.rel_pos_transformer = RelativePositionTransformer(in_channels, out_channels, hidden_channels, **params)
+
+    def forward(self, x, x_mask=None, g=None):  # pylint: disable=unused-argument
+        o = self.prenet(x) * x_mask
+        o = self.rel_pos_transformer(o, x_mask)
+        return o
+
+
+class FFTransformerDecoder(nn.Module):
+    """Decoder with FeedForwardTransformer.
+
+    Default params
+            params={
+                'hidden_channels_ffn': 1024,
+                'num_heads': 2,
+                "dropout_p": 0.1,
+                "num_layers": 6,
+            }
+
+    Args:
+        in_channels (int): number of input channels.
+        out_channels (int): number of output channels.
+        hidden_channels (int): number of hidden channels including Transformer layers.
+        params (dict): dictionary for residual convolutional blocks.
+    """
+
+    def __init__(self, in_channels, out_channels, params):
+        super().__init__()
+        self.transformer_block = FFTransformerBlock(in_channels, **params)
+        self.postnet = nn.Conv1d(in_channels, out_channels, 1)
+
+    def forward(self, x, x_mask=None, g=None):  # pylint: disable=unused-argument
+        # TODO: handle multi-speaker
+        x_mask = 1 if x_mask is None else x_mask
+        o = self.transformer_block(x) * x_mask
+        o = self.postnet(o) * x_mask
+        return o
+
+
+class ResidualConv1dBNDecoder(nn.Module):
+    """Residual Convolutional Decoder as in the original Speedy Speech paper
+
+    TODO: Integrate speaker conditioning vector.
+
+    Note:
+        Default params
+                params = {
+                    "kernel_size": 4,
+                    "dilations": 4 * [1, 2, 4, 8] + [1],
+                    "num_conv_blocks": 2,
+                    "num_res_blocks": 17
+                }
+
+    Args:
+        in_channels (int): number of input channels.
+        out_channels (int): number of output channels.
+        hidden_channels (int): number of hidden channels including ResidualConv1dBNBlock layers.
+        params (dict): dictionary for residual convolutional blocks.
+    """
+
+    def __init__(self, in_channels, out_channels, hidden_channels, params):
+        super().__init__()
+        self.res_conv_block = ResidualConv1dBNBlock(in_channels, hidden_channels, hidden_channels, **params)
+        self.post_conv = nn.Conv1d(hidden_channels, hidden_channels, 1)
+        self.postnet = nn.Sequential(
+            Conv1dBNBlock(
+                hidden_channels, hidden_channels, hidden_channels, params["kernel_size"], 1, num_conv_blocks=2
+            ),
+            nn.Conv1d(hidden_channels, out_channels, 1),
+        )
+
+    def forward(self, x, x_mask=None, g=None):  # pylint: disable=unused-argument
+        o = self.res_conv_block(x, x_mask)
+        o = self.post_conv(o) + x
+        return self.postnet(o) * x_mask
+
+
+class Decoder(nn.Module):
+    """Decodes the expanded phoneme encoding into spectrograms
+    Args:
+        out_channels (int): number of output channels.
+        in_hidden_channels (int): input and hidden channels. Model keeps the input channels for the intermediate layers.
+        decoder_type (str): decoder layer types. 'transformers' or 'residual_conv_bn'. Default 'residual_conv_bn'.
+        decoder_params (dict): model parameters for specified decoder type.
+        c_in_channels (int): number of channels for conditional input.
+
+    Shapes:
+        - input: (B, C, T)
+    """
+
+    # pylint: disable=dangerous-default-value
+    def __init__(
+        self,
+        out_channels,
+        in_hidden_channels,
+        decoder_type="residual_conv_bn",
+        decoder_params={
+            "kernel_size": 4,
+            "dilations": 4 * [1, 2, 4, 8] + [1],
+            "num_conv_blocks": 2,
+            "num_res_blocks": 17,
+        },
+        c_in_channels=0,
+    ):
+        super().__init__()
+
+        if decoder_type.lower() == "relative_position_transformer":
+            self.decoder = RelativePositionTransformerDecoder(
+                in_channels=in_hidden_channels,
+                out_channels=out_channels,
+                hidden_channels=in_hidden_channels,
+                params=decoder_params,
+            )
+        elif decoder_type.lower() == "residual_conv_bn":
+            self.decoder = ResidualConv1dBNDecoder(
+                in_channels=in_hidden_channels,
+                out_channels=out_channels,
+                hidden_channels=in_hidden_channels,
+                params=decoder_params,
+            )
+        elif decoder_type.lower() == "wavenet":
+            self.decoder = WaveNetDecoder(
+                in_channels=in_hidden_channels,
+                out_channels=out_channels,
+                hidden_channels=in_hidden_channels,
+                c_in_channels=c_in_channels,
+                params=decoder_params,
+            )
+        elif decoder_type.lower() == "fftransformer":
+            self.decoder = FFTransformerDecoder(in_hidden_channels, out_channels, decoder_params)
+        else:
+            raise ValueError(f"[!] Unknown decoder type - {decoder_type}")
+
+    def forward(self, x, x_mask, g=None):  # pylint: disable=unused-argument
+        """
+        Args:
+            x: [B, C, T]
+            x_mask: [B, 1, T]
+            g: [B, C_g, 1]
+        """
+        # TODO: implement multi-speaker
+        o = self.decoder(x, x_mask, g)
+        return o

TTS/tts/layers/feed_forward/duration_predictor.py

@@ -0,0 +1,41 @@
+from torch import nn
+
+from TTS.tts.layers.generic.res_conv_bn import Conv1dBN
+
+
+class DurationPredictor(nn.Module):
+    """Speedy Speech duration predictor model.
+    Predicts phoneme durations from encoder outputs.
+
+    Note:
+        Outputs interpreted as log(durations)
+        To get actual durations, do exp transformation
+
+    conv_BN_4x1 -> conv_BN_3x1 -> conv_BN_1x1 -> conv_1x1
+
+    Args:
+        hidden_channels (int): number of channels in the inner layers.
+    """
+
+    def __init__(self, hidden_channels):
+        super().__init__()
+
+        self.layers = nn.ModuleList(
+            [
+                Conv1dBN(hidden_channels, hidden_channels, 4, 1),
+                Conv1dBN(hidden_channels, hidden_channels, 3, 1),
+                Conv1dBN(hidden_channels, hidden_channels, 1, 1),
+                nn.Conv1d(hidden_channels, 1, 1),
+            ]
+        )
+
+    def forward(self, x, x_mask):
+        """
+        Shapes:
+            x: [B, C, T]
+            x_mask: [B, 1, T]
+        """
+        o = x
+        for layer in self.layers:
+            o = layer(o) * x_mask
+        return o

TTS/tts/layers/feed_forward/encoder.py

@@ -0,0 +1,162 @@
+from torch import nn
+
+from TTS.tts.layers.generic.res_conv_bn import ResidualConv1dBNBlock
+from TTS.tts.layers.generic.transformer import FFTransformerBlock
+from TTS.tts.layers.glow_tts.transformer import RelativePositionTransformer
+
+
+class RelativePositionTransformerEncoder(nn.Module):
+    """Speedy speech encoder built on Transformer with Relative Position encoding.
+
+    TODO: Integrate speaker conditioning vector.
+
+    Args:
+        in_channels (int): number of input channels.
+        out_channels (int): number of output channels.
+        hidden_channels (int): number of hidden channels
+        params (dict): dictionary for residual convolutional blocks.
+    """
+
+    def __init__(self, in_channels, out_channels, hidden_channels, params):
+        super().__init__()
+        self.prenet = ResidualConv1dBNBlock(
+            in_channels,
+            hidden_channels,
+            hidden_channels,
+            kernel_size=5,
+            num_res_blocks=3,
+            num_conv_blocks=1,
+            dilations=[1, 1, 1],
+        )
+        self.rel_pos_transformer = RelativePositionTransformer(hidden_channels, out_channels, hidden_channels, **params)
+
+    def forward(self, x, x_mask=None, g=None):  # pylint: disable=unused-argument
+        if x_mask is None:
+            x_mask = 1
+        o = self.prenet(x) * x_mask
+        o = self.rel_pos_transformer(o, x_mask)
+        return o
+
+
+class ResidualConv1dBNEncoder(nn.Module):
+    """Residual Convolutional Encoder as in the original Speedy Speech paper
+
+    TODO: Integrate speaker conditioning vector.
+
+    Args:
+        in_channels (int): number of input channels.
+        out_channels (int): number of output channels.
+        hidden_channels (int): number of hidden channels
+        params (dict): dictionary for residual convolutional blocks.
+    """
+
+    def __init__(self, in_channels, out_channels, hidden_channels, params):
+        super().__init__()
+        self.prenet = nn.Sequential(nn.Conv1d(in_channels, hidden_channels, 1), nn.ReLU())
+        self.res_conv_block = ResidualConv1dBNBlock(hidden_channels, hidden_channels, hidden_channels, **params)
+
+        self.postnet = nn.Sequential(
+            *[
+                nn.Conv1d(hidden_channels, hidden_channels, 1),
+                nn.ReLU(),
+                nn.BatchNorm1d(hidden_channels),
+                nn.Conv1d(hidden_channels, out_channels, 1),
+            ]
+        )
+
+    def forward(self, x, x_mask=None, g=None):  # pylint: disable=unused-argument
+        if x_mask is None:
+            x_mask = 1
+        o = self.prenet(x) * x_mask
+        o = self.res_conv_block(o, x_mask)
+        o = self.postnet(o + x) * x_mask
+        return o * x_mask
+
+
+class Encoder(nn.Module):
+    # pylint: disable=dangerous-default-value
+    """Factory class for Speedy Speech encoder enables different encoder types internally.
+
+    Args:
+        num_chars (int): number of characters.
+        out_channels (int): number of output channels.
+        in_hidden_channels (int): input and hidden channels. Model keeps the input channels for the intermediate layers.
+        encoder_type (str): encoder layer types. 'transformers' or 'residual_conv_bn'. Default 'residual_conv_bn'.
+        encoder_params (dict): model parameters for specified encoder type.
+        c_in_channels (int): number of channels for conditional input.
+
+    Note:
+        Default encoder_params to be set in config.json...
+
+        ```python
+        # for 'relative_position_transformer'
+        encoder_params={
+            'hidden_channels_ffn': 128,
+            'num_heads': 2,
+            "kernel_size": 3,
+            "dropout_p": 0.1,
+            "num_layers": 6,
+            "rel_attn_window_size": 4,
+            "input_length": None
+        },
+
+        # for 'residual_conv_bn'
+        encoder_params = {
+            "kernel_size": 4,
+            "dilations": 4 * [1, 2, 4] + [1],
+            "num_conv_blocks": 2,
+            "num_res_blocks": 13
+        }
+
+        # for 'fftransformer'
+        encoder_params = {
+            "hidden_channels_ffn": 1024 ,
+            "num_heads": 2,
+            "num_layers": 6,
+            "dropout_p": 0.1
+        }
+        ```
+    """
+
+    def __init__(
+        self,
+        in_hidden_channels,
+        out_channels,
+        encoder_type="residual_conv_bn",
+        encoder_params={"kernel_size": 4, "dilations": 4 * [1, 2, 4] + [1], "num_conv_blocks": 2, "num_res_blocks": 13},
+        c_in_channels=0,
+    ):
+        super().__init__()
+        self.out_channels = out_channels
+        self.in_channels = in_hidden_channels
+        self.hidden_channels = in_hidden_channels
+        self.encoder_type = encoder_type
+        self.c_in_channels = c_in_channels
+
+        # init encoder
+        if encoder_type.lower() == "relative_position_transformer":
+            # text encoder
+            # pylint: disable=unexpected-keyword-arg
+            self.encoder = RelativePositionTransformerEncoder(
+                in_hidden_channels, out_channels, in_hidden_channels, encoder_params
+            )
+        elif encoder_type.lower() == "residual_conv_bn":
+            self.encoder = ResidualConv1dBNEncoder(in_hidden_channels, out_channels, in_hidden_channels, encoder_params)
+        elif encoder_type.lower() == "fftransformer":
+            assert (
+                in_hidden_channels == out_channels
+            ), "[!] must be `in_channels` == `out_channels` when encoder type is 'fftransformer'"
+            # pylint: disable=unexpected-keyword-arg
+            self.encoder = FFTransformerBlock(in_hidden_channels, **encoder_params)
+        else:
+            raise NotImplementedError(" [!] unknown encoder type.")
+
+    def forward(self, x, x_mask, g=None):  # pylint: disable=unused-argument
+        """
+        Shapes:
+            x: [B, C, T]
+            x_mask: [B, 1, T]
+            g: [B, C, 1]
+        """
+        o = self.encoder(x, x_mask)
+        return o * x_mask

TTS/tts/layers/generic/aligner.py

@@ -0,0 +1,92 @@
+from typing import Tuple
+
+import torch
+from torch import nn
+
+
+class AlignmentNetwork(torch.nn.Module):
+    """Aligner Network for learning alignment between the input text and the model output with Gaussian Attention.
+
+    ::
+
+        query -> conv1d -> relu -> conv1d -> relu -> conv1d -> L2_dist -> softmax -> alignment
+        key   -> conv1d -> relu -> conv1d -----------------------^
+
+    Args:
+        in_query_channels (int): Number of channels in the query network. Defaults to 80.
+        in_key_channels (int): Number of channels in the key network. Defaults to 512.
+        attn_channels (int): Number of inner channels in the attention layers. Defaults to 80.
+        temperature (float): Temperature for the softmax. Defaults to 0.0005.
+    """
+
+    def __init__(
+        self,
+        in_query_channels=80,
+        in_key_channels=512,
+        attn_channels=80,
+        temperature=0.0005,
+    ):
+        super().__init__()
+        self.temperature = temperature
+        self.softmax = torch.nn.Softmax(dim=3)
+        self.log_softmax = torch.nn.LogSoftmax(dim=3)
+
+        self.key_layer = nn.Sequential(
+            nn.Conv1d(
+                in_key_channels,
+                in_key_channels * 2,
+                kernel_size=3,
+                padding=1,
+                bias=True,
+            ),
+            torch.nn.ReLU(),
+            nn.Conv1d(in_key_channels * 2, attn_channels, kernel_size=1, padding=0, bias=True),
+        )
+
+        self.query_layer = nn.Sequential(
+            nn.Conv1d(
+                in_query_channels,
+                in_query_channels * 2,
+                kernel_size=3,
+                padding=1,
+                bias=True,
+            ),
+            torch.nn.ReLU(),
+            nn.Conv1d(in_query_channels * 2, in_query_channels, kernel_size=1, padding=0, bias=True),
+            torch.nn.ReLU(),
+            nn.Conv1d(in_query_channels, attn_channels, kernel_size=1, padding=0, bias=True),
+        )
+
+        self.init_layers()
+
+    def init_layers(self):
+        torch.nn.init.xavier_uniform_(self.key_layer[0].weight, gain=torch.nn.init.calculate_gain("relu"))
+        torch.nn.init.xavier_uniform_(self.key_layer[2].weight, gain=torch.nn.init.calculate_gain("linear"))
+        torch.nn.init.xavier_uniform_(self.query_layer[0].weight, gain=torch.nn.init.calculate_gain("relu"))
+        torch.nn.init.xavier_uniform_(self.query_layer[2].weight, gain=torch.nn.init.calculate_gain("linear"))
+        torch.nn.init.xavier_uniform_(self.query_layer[4].weight, gain=torch.nn.init.calculate_gain("linear"))
+
+    def forward(
+        self, queries: torch.tensor, keys: torch.tensor, mask: torch.tensor = None, attn_prior: torch.tensor = None
+    ) -> Tuple[torch.tensor, torch.tensor]:
+        """Forward pass of the aligner encoder.
+        Shapes:
+            - queries: :math:`[B, C, T_de]`
+            - keys: :math:`[B, C_emb, T_en]`
+            - mask: :math:`[B, T_de]`
+        Output:
+            attn (torch.tensor): :math:`[B, 1, T_en, T_de]` soft attention mask.
+            attn_logp (torch.tensor): :math:`[ßB, 1, T_en , T_de]` log probabilities.
+        """
+        key_out = self.key_layer(keys)
+        query_out = self.query_layer(queries)
+        attn_factor = (query_out[:, :, :, None] - key_out[:, :, None]) ** 2
+        attn_logp = -self.temperature * attn_factor.sum(1, keepdim=True)
+        if attn_prior is not None:
+            attn_logp = self.log_softmax(attn_logp) + torch.log(attn_prior[:, None] + 1e-8)
+
+        if mask is not None:
+            attn_logp.data.masked_fill_(~mask.bool().unsqueeze(2), -float("inf"))
+
+        attn = self.softmax(attn_logp)
+        return attn, attn_logp

TTS/tts/layers/generic/gated_conv.py

@@ -0,0 +1,37 @@
+from torch import nn
+
+from .normalization import LayerNorm
+
+
+class GatedConvBlock(nn.Module):
+    """Gated convolutional block as in https://arxiv.org/pdf/1612.08083.pdf
+    Args:
+        in_out_channels (int): number of input/output channels.
+        kernel_size (int): convolution kernel size.
+        dropout_p (float): dropout rate.
+    """
+
+    def __init__(self, in_out_channels, kernel_size, dropout_p, num_layers):
+        super().__init__()
+        # class arguments
+        self.dropout_p = dropout_p
+        self.num_layers = num_layers
+        # define layers
+        self.conv_layers = nn.ModuleList()
+        self.norm_layers = nn.ModuleList()
+        self.layers = nn.ModuleList()
+        for _ in range(num_layers):
+            self.conv_layers += [nn.Conv1d(in_out_channels, 2 * in_out_channels, kernel_size, padding=kernel_size // 2)]
+            self.norm_layers += [LayerNorm(2 * in_out_channels)]
+
+    def forward(self, x, x_mask):
+        o = x
+        res = x
+        for idx in range(self.num_layers):
+            o = nn.functional.dropout(o, p=self.dropout_p, training=self.training)
+            o = self.conv_layers[idx](o * x_mask)
+            o = self.norm_layers[idx](o)
+            o = nn.functional.glu(o, dim=1)
+            o = res + o
+            res = o
+        return o

TTS/tts/layers/generic/normalization.py

@@ -0,0 +1,123 @@
+import torch
+from torch import nn
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, channels, eps=1e-4):
+        """Layer norm for the 2nd dimension of the input.
+        Args:
+            channels (int): number of channels (2nd dimension) of the input.
+            eps (float): to prevent 0 division
+
+        Shapes:
+            - input: (B, C, T)
+            - output: (B, C, T)
+        """
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(1, channels, 1) * 0.1)
+        self.beta = nn.Parameter(torch.zeros(1, channels, 1))
+
+    def forward(self, x):
+        mean = torch.mean(x, 1, keepdim=True)
+        variance = torch.mean((x - mean) ** 2, 1, keepdim=True)
+        x = (x - mean) * torch.rsqrt(variance + self.eps)
+        x = x * self.gamma + self.beta
+        return x
+
+
+class LayerNorm2(nn.Module):
+    """Layer norm for the 2nd dimension of the input using torch primitive.
+    Args:
+        channels (int): number of channels (2nd dimension) of the input.
+        eps (float): to prevent 0 division
+
+    Shapes:
+        - input: (B, C, T)
+        - output: (B, C, T)
+    """
+
+    def __init__(self, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        x = x.transpose(1, -1)
+        x = torch.nn.functional.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+class TemporalBatchNorm1d(nn.BatchNorm1d):
+    """Normalize each channel separately over time and batch."""
+
+    def __init__(self, channels, affine=True, track_running_stats=True, momentum=0.1):
+        super().__init__(channels, affine=affine, track_running_stats=track_running_stats, momentum=momentum)
+
+    def forward(self, x):
+        return super().forward(x.transpose(2, 1)).transpose(2, 1)
+
+
+class ActNorm(nn.Module):
+    """Activation Normalization bijector as an alternative to Batch Norm. It computes
+    mean and std from a sample data in advance and it uses these values
+    for normalization at training.
+
+    Args:
+        channels (int): input channels.
+        ddi (False): data depended initialization flag.
+
+    Shapes:
+        - inputs: (B, C, T)
+        - outputs: (B, C, T)
+    """
+
+    def __init__(self, channels, ddi=False, **kwargs):  # pylint: disable=unused-argument
+        super().__init__()
+        self.channels = channels
+        self.initialized = not ddi
+
+        self.logs = nn.Parameter(torch.zeros(1, channels, 1))
+        self.bias = nn.Parameter(torch.zeros(1, channels, 1))
+
+    def forward(self, x, x_mask=None, reverse=False, **kwargs):  # pylint: disable=unused-argument
+        if x_mask is None:
+            x_mask = torch.ones(x.size(0), 1, x.size(2)).to(device=x.device, dtype=x.dtype)
+        x_len = torch.sum(x_mask, [1, 2])
+        if not self.initialized:
+            self.initialize(x, x_mask)
+            self.initialized = True
+
+        if reverse:
+            z = (x - self.bias) * torch.exp(-self.logs) * x_mask
+            logdet = None
+        else:
+            z = (self.bias + torch.exp(self.logs) * x) * x_mask
+            logdet = torch.sum(self.logs) * x_len  # [b]
+
+        return z, logdet
+
+    def store_inverse(self):
+        pass
+
+    def set_ddi(self, ddi):
+        self.initialized = not ddi
+
+    def initialize(self, x, x_mask):
+        with torch.no_grad():
+            denom = torch.sum(x_mask, [0, 2])
+            m = torch.sum(x * x_mask, [0, 2]) / denom
+            m_sq = torch.sum(x * x * x_mask, [0, 2]) / denom
+            v = m_sq - (m**2)
+            logs = 0.5 * torch.log(torch.clamp_min(v, 1e-6))
+
+            bias_init = (-m * torch.exp(-logs)).view(*self.bias.shape).to(dtype=self.bias.dtype)
+            logs_init = (-logs).view(*self.logs.shape).to(dtype=self.logs.dtype)
+
+            self.bias.data.copy_(bias_init)
+            self.logs.data.copy_(logs_init)

TTS/tts/layers/generic/pos_encoding.py

@@ -0,0 +1,69 @@
+import math
+
+import torch
+from torch import nn
+
+
+class PositionalEncoding(nn.Module):
+    """Sinusoidal positional encoding for non-recurrent neural networks.
+    Implementation based on "Attention Is All You Need"
+
+    Args:
+       channels (int): embedding size
+       dropout_p (float): dropout rate applied to the output.
+       max_len (int): maximum sequence length.
+       use_scale (bool): whether to use a learnable scaling coefficient.
+    """
+
+    def __init__(self, channels, dropout_p=0.0, max_len=5000, use_scale=False):
+        super().__init__()
+        if channels % 2 != 0:
+            raise ValueError(
+                "Cannot use sin/cos positional encoding with " "odd channels (got channels={:d})".format(channels)
+            )
+        self.use_scale = use_scale
+        if use_scale:
+            self.scale = torch.nn.Parameter(torch.ones(1))
+        pe = torch.zeros(max_len, channels)
+        position = torch.arange(0, max_len).unsqueeze(1)
+        div_term = torch.pow(10000, torch.arange(0, channels, 2).float() / channels)
+        pe[:, 0::2] = torch.sin(position.float() * div_term)
+        pe[:, 1::2] = torch.cos(position.float() * div_term)
+        pe = pe.unsqueeze(0).transpose(1, 2)
+        self.register_buffer("pe", pe)
+        if dropout_p > 0:
+            self.dropout = nn.Dropout(p=dropout_p)
+        self.channels = channels
+
+    def forward(self, x, mask=None, first_idx=None, last_idx=None):
+        """
+        Shapes:
+            x: [B, C, T]
+            mask: [B, 1, T]
+            first_idx: int
+            last_idx: int
+        """
+
+        x = x * math.sqrt(self.channels)
+        if first_idx is None:
+            if self.pe.size(2) < x.size(2):
+                raise RuntimeError(
+                    f"Sequence is {x.size(2)} but PositionalEncoding is"
+                    f" limited to {self.pe.size(2)}. See max_len argument."
+                )
+            if mask is not None:
+                pos_enc = self.pe[:, :, : x.size(2)] * mask
+            else:
+                pos_enc = self.pe[:, :, : x.size(2)]
+            if self.use_scale:
+                x = x + self.scale * pos_enc
+            else:
+                x = x + pos_enc
+        else:
+            if self.use_scale:
+                x = x + self.scale * self.pe[:, :, first_idx:last_idx]
+            else:
+                x = x + self.pe[:, :, first_idx:last_idx]
+        if hasattr(self, "dropout"):
+            x = self.dropout(x)
+        return x

TTS/tts/layers/generic/res_conv_bn.py

@@ -0,0 +1,127 @@
+from torch import nn
+
+
+class ZeroTemporalPad(nn.Module):
+    """Pad sequences to equal lentgh in the temporal dimension"""
+
+    def __init__(self, kernel_size, dilation):
+        super().__init__()
+        total_pad = dilation * (kernel_size - 1)
+        begin = total_pad // 2
+        end = total_pad - begin
+        self.pad_layer = nn.ZeroPad2d((0, 0, begin, end))
+
+    def forward(self, x):
+        return self.pad_layer(x)
+
+
+class Conv1dBN(nn.Module):
+    """1d convolutional with batch norm.
+    conv1d -> relu -> BN blocks.
+
+    Note:
+        Batch normalization is applied after ReLU regarding the original implementation.
+
+    Args:
+        in_channels (int): number of input channels.
+        out_channels (int): number of output channels.
+        kernel_size (int): kernel size for convolutional filters.
+        dilation (int): dilation for convolution layers.
+    """
+
+    def __init__(self, in_channels, out_channels, kernel_size, dilation):
+        super().__init__()
+        padding = dilation * (kernel_size - 1)
+        pad_s = padding // 2
+        pad_e = padding - pad_s
+        self.conv1d = nn.Conv1d(in_channels, out_channels, kernel_size, dilation=dilation)
+        self.pad = nn.ZeroPad2d((pad_s, pad_e, 0, 0))  # uneven left and right padding
+        self.norm = nn.BatchNorm1d(out_channels)
+
+    def forward(self, x):
+        o = self.conv1d(x)
+        o = self.pad(o)
+        o = nn.functional.relu(o)
+        o = self.norm(o)
+        return o
+
+
+class Conv1dBNBlock(nn.Module):
+    """1d convolutional block with batch norm. It is a set of conv1d -> relu -> BN blocks.
+
+    Args:
+        in_channels (int): number of input channels.
+        out_channels (int): number of output channels.
+        hidden_channels (int): number of inner convolution channels.
+        kernel_size (int): kernel size for convolutional filters.
+        dilation (int): dilation for convolution layers.
+        num_conv_blocks (int, optional): number of convolutional blocks. Defaults to 2.
+    """
+
+    def __init__(self, in_channels, out_channels, hidden_channels, kernel_size, dilation, num_conv_blocks=2):
+        super().__init__()
+        self.conv_bn_blocks = []
+        for idx in range(num_conv_blocks):
+            layer = Conv1dBN(
+                in_channels if idx == 0 else hidden_channels,
+                out_channels if idx == (num_conv_blocks - 1) else hidden_channels,
+                kernel_size,
+                dilation,
+            )
+            self.conv_bn_blocks.append(layer)
+        self.conv_bn_blocks = nn.Sequential(*self.conv_bn_blocks)
+
+    def forward(self, x):
+        """
+        Shapes:
+            x: (B, D, T)
+        """
+        return self.conv_bn_blocks(x)
+
+
+class ResidualConv1dBNBlock(nn.Module):
+    """Residual Convolutional Blocks with BN
+    Each block has 'num_conv_block' conv layers and 'num_res_blocks' such blocks are connected
+    with residual connections.
+
+    conv_block = (conv1d -> relu -> bn) x 'num_conv_blocks'
+    residuak_conv_block =  (x -> conv_block ->  + ->) x 'num_res_blocks'
+                            ' - - - - - - - - - ^
+    Args:
+        in_channels (int): number of input channels.
+        out_channels (int): number of output channels.
+        hidden_channels (int): number of inner convolution channels.
+        kernel_size (int): kernel size for convolutional filters.
+        dilations (list): dilations for each convolution layer.
+        num_res_blocks (int, optional): number of residual blocks. Defaults to 13.
+        num_conv_blocks (int, optional): number of convolutional blocks in each residual block. Defaults to 2.
+    """
+
+    def __init__(
+        self, in_channels, out_channels, hidden_channels, kernel_size, dilations, num_res_blocks=13, num_conv_blocks=2
+    ):
+        super().__init__()
+        assert len(dilations) == num_res_blocks
+        self.res_blocks = nn.ModuleList()
+        for idx, dilation in enumerate(dilations):
+            block = Conv1dBNBlock(
+                in_channels if idx == 0 else hidden_channels,
+                out_channels if (idx + 1) == len(dilations) else hidden_channels,
+                hidden_channels,
+                kernel_size,
+                dilation,
+                num_conv_blocks,
+            )
+            self.res_blocks.append(block)
+
+    def forward(self, x, x_mask=None):
+        if x_mask is None:
+            x_mask = 1.0
+        o = x * x_mask
+        for block in self.res_blocks:
+            res = o
+            o = block(o)
+            o = o + res
+            if x_mask is not None:
+                o = o * x_mask
+        return o

TTS/tts/layers/generic/time_depth_sep_conv.py

@@ -0,0 +1,84 @@
+import torch
+from torch import nn
+
+
+class TimeDepthSeparableConv(nn.Module):
+    """Time depth separable convolution as in https://arxiv.org/pdf/1904.02619.pdf
+    It shows competative results with less computation and memory footprint."""
+
+    def __init__(self, in_channels, hid_channels, out_channels, kernel_size, bias=True):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.hid_channels = hid_channels
+        self.kernel_size = kernel_size
+
+        self.time_conv = nn.Conv1d(
+            in_channels,
+            2 * hid_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=bias,
+        )
+        self.norm1 = nn.BatchNorm1d(2 * hid_channels)
+        self.depth_conv = nn.Conv1d(
+            hid_channels,
+            hid_channels,
+            kernel_size,
+            stride=1,
+            padding=(kernel_size - 1) // 2,
+            groups=hid_channels,
+            bias=bias,
+        )
+        self.norm2 = nn.BatchNorm1d(hid_channels)
+        self.time_conv2 = nn.Conv1d(
+            hid_channels,
+            out_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=bias,
+        )
+        self.norm3 = nn.BatchNorm1d(out_channels)
+
+    def forward(self, x):
+        x_res = x
+        x = self.time_conv(x)
+        x = self.norm1(x)
+        x = nn.functional.glu(x, dim=1)
+        x = self.depth_conv(x)
+        x = self.norm2(x)
+        x = x * torch.sigmoid(x)
+        x = self.time_conv2(x)
+        x = self.norm3(x)
+        x = x_res + x
+        return x
+
+
+class TimeDepthSeparableConvBlock(nn.Module):
+    def __init__(self, in_channels, hid_channels, out_channels, num_layers, kernel_size, bias=True):
+        super().__init__()
+        assert (kernel_size - 1) % 2 == 0
+        assert num_layers > 1
+
+        self.layers = nn.ModuleList()
+        layer = TimeDepthSeparableConv(
+            in_channels, hid_channels, out_channels if num_layers == 1 else hid_channels, kernel_size, bias
+        )
+        self.layers.append(layer)
+        for idx in range(num_layers - 1):
+            layer = TimeDepthSeparableConv(
+                hid_channels,
+                hid_channels,
+                out_channels if (idx + 1) == (num_layers - 1) else hid_channels,
+                kernel_size,
+                bias,
+            )
+            self.layers.append(layer)
+
+    def forward(self, x, mask):
+        for layer in self.layers:
+            x = layer(x * mask)
+        return x

TTS/tts/layers/generic/transformer.py

@@ -0,0 +1,89 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class FFTransformer(nn.Module):
+    def __init__(self, in_out_channels, num_heads, hidden_channels_ffn=1024, kernel_size_fft=3, dropout_p=0.1):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(in_out_channels, num_heads, dropout=dropout_p)
+
+        padding = (kernel_size_fft - 1) // 2
+        self.conv1 = nn.Conv1d(in_out_channels, hidden_channels_ffn, kernel_size=kernel_size_fft, padding=padding)
+        self.conv2 = nn.Conv1d(hidden_channels_ffn, in_out_channels, kernel_size=kernel_size_fft, padding=padding)
+
+        self.norm1 = nn.LayerNorm(in_out_channels)
+        self.norm2 = nn.LayerNorm(in_out_channels)
+
+        self.dropout1 = nn.Dropout(dropout_p)
+        self.dropout2 = nn.Dropout(dropout_p)
+
+    def forward(self, src, src_mask=None, src_key_padding_mask=None):
+        """😦 ugly looking with all the transposing"""
+        src = src.permute(2, 0, 1)
+        src2, enc_align = self.self_attn(src, src, src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)
+        src = src + self.dropout1(src2)
+        src = self.norm1(src + src2)
+        # T x B x D -> B x D x T
+        src = src.permute(1, 2, 0)
+        src2 = self.conv2(F.relu(self.conv1(src)))
+        src2 = self.dropout2(src2)
+        src = src + src2
+        src = src.transpose(1, 2)
+        src = self.norm2(src)
+        src = src.transpose(1, 2)
+        return src, enc_align
+
+
+class FFTransformerBlock(nn.Module):
+    def __init__(self, in_out_channels, num_heads, hidden_channels_ffn, num_layers, dropout_p):
+        super().__init__()
+        self.fft_layers = nn.ModuleList(
+            [
+                FFTransformer(
+                    in_out_channels=in_out_channels,
+                    num_heads=num_heads,
+                    hidden_channels_ffn=hidden_channels_ffn,
+                    dropout_p=dropout_p,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+    def forward(self, x, mask=None, g=None):  # pylint: disable=unused-argument
+        """
+        TODO: handle multi-speaker
+        Shapes:
+            - x: :math:`[B, C, T]`
+            - mask:  :math:`[B, 1, T] or [B, T]`
+        """
+        if mask is not None and mask.ndim == 3:
+            mask = mask.squeeze(1)
+            # mask is negated, torch uses 1s and 0s reversely.
+            mask = ~mask.bool()
+        alignments = []
+        for layer in self.fft_layers:
+            x, align = layer(x, src_key_padding_mask=mask)
+            alignments.append(align.unsqueeze(1))
+        alignments = torch.cat(alignments, 1)
+        return x
+
+
+class FFTDurationPredictor:
+    def __init__(
+        self, in_channels, hidden_channels, num_heads, num_layers, dropout_p=0.1, cond_channels=None
+    ):  # pylint: disable=unused-argument
+        self.fft = FFTransformerBlock(in_channels, num_heads, hidden_channels, num_layers, dropout_p)
+        self.proj = nn.Linear(in_channels, 1)
+
+    def forward(self, x, mask=None, g=None):  # pylint: disable=unused-argument
+        """
+        Shapes:
+            - x: :math:`[B, C, T]`
+            - mask:  :math:`[B, 1, T]`
+
+        TODO: Handle the cond input
+        """
+        x = self.fft(x, mask=mask)
+        x = self.proj(x)
+        return x

TTS/tts/layers/generic/wavenet.py

@@ -0,0 +1,176 @@
+import torch
+from torch import nn
+from torch.nn.utils import parametrize
+
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
+
+
+class WN(torch.nn.Module):
+    """Wavenet layers with weight norm and no input conditioning.
+
+         |-----------------------------------------------------------------------------|
+         |                                    |-> tanh    -|                           |
+    res -|- conv1d(dilation) -> dropout -> + -|            * -> conv1d1x1 -> split -|- + -> res
+    g -------------------------------------|  |-> sigmoid -|                        |
+    o --------------------------------------------------------------------------- + --------- o
+
+    Args:
+        in_channels (int): number of input channels.
+        hidden_channes (int): number of hidden channels.
+        kernel_size (int): filter kernel size for the first conv layer.
+        dilation_rate (int): dilations rate to increase dilation per layer.
+            If it is 2, dilations are 1, 2, 4, 8 for the next 4 layers.
+        num_layers (int): number of wavenet layers.
+        c_in_channels (int): number of channels of conditioning input.
+        dropout_p (float): dropout rate.
+        weight_norm (bool): enable/disable weight norm for convolution layers.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        num_layers,
+        c_in_channels=0,
+        dropout_p=0,
+        weight_norm=True,
+    ):
+        super().__init__()
+        assert kernel_size % 2 == 1
+        assert hidden_channels % 2 == 0
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.num_layers = num_layers
+        self.c_in_channels = c_in_channels
+        self.dropout_p = dropout_p
+
+        self.in_layers = torch.nn.ModuleList()
+        self.res_skip_layers = torch.nn.ModuleList()
+        self.dropout = nn.Dropout(dropout_p)
+
+        # init conditioning layer
+        if c_in_channels > 0:
+            cond_layer = torch.nn.Conv1d(c_in_channels, 2 * hidden_channels * num_layers, 1)
+            self.cond_layer = torch.nn.utils.parametrizations.weight_norm(cond_layer, name="weight")
+        # intermediate layers
+        for i in range(num_layers):
+            dilation = dilation_rate**i
+            padding = int((kernel_size * dilation - dilation) / 2)
+            if i == 0:
+                in_layer = torch.nn.Conv1d(
+                    in_channels, 2 * hidden_channels, kernel_size, dilation=dilation, padding=padding
+                )
+            else:
+                in_layer = torch.nn.Conv1d(
+                    hidden_channels, 2 * hidden_channels, kernel_size, dilation=dilation, padding=padding
+                )
+            in_layer = torch.nn.utils.parametrizations.weight_norm(in_layer, name="weight")
+            self.in_layers.append(in_layer)
+
+            if i < num_layers - 1:
+                res_skip_channels = 2 * hidden_channels
+            else:
+                res_skip_channels = hidden_channels
+
+            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
+            res_skip_layer = torch.nn.utils.parametrizations.weight_norm(res_skip_layer, name="weight")
+            self.res_skip_layers.append(res_skip_layer)
+        # setup weight norm
+        if not weight_norm:
+            self.remove_weight_norm()
+
+    def forward(self, x, x_mask=None, g=None, **kwargs):  # pylint: disable=unused-argument
+        output = torch.zeros_like(x)
+        n_channels_tensor = torch.IntTensor([self.hidden_channels])
+        x_mask = 1.0 if x_mask is None else x_mask
+        if g is not None:
+            g = self.cond_layer(g)
+        for i in range(self.num_layers):
+            x_in = self.in_layers[i](x)
+            x_in = self.dropout(x_in)
+            if g is not None:
+                cond_offset = i * 2 * self.hidden_channels
+                g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
+            else:
+                g_l = torch.zeros_like(x_in)
+            acts = fused_add_tanh_sigmoid_multiply(x_in, g_l, n_channels_tensor)
+            res_skip_acts = self.res_skip_layers[i](acts)
+            if i < self.num_layers - 1:
+                x = (x + res_skip_acts[:, : self.hidden_channels, :]) * x_mask
+                output = output + res_skip_acts[:, self.hidden_channels :, :]
+            else:
+                output = output + res_skip_acts
+        return output * x_mask
+
+    def remove_weight_norm(self):
+        if self.c_in_channels != 0:
+            parametrize.remove_parametrizations(self.cond_layer, "weight")
+        for l in self.in_layers:
+            parametrize.remove_parametrizations(l, "weight")
+        for l in self.res_skip_layers:
+            parametrize.remove_parametrizations(l, "weight")
+
+
+class WNBlocks(nn.Module):
+    """Wavenet blocks.
+
+    Note: After each block dilation resets to 1 and it increases in each block
+        along the dilation rate.
+
+    Args:
+        in_channels (int): number of input channels.
+        hidden_channes (int): number of hidden channels.
+        kernel_size (int): filter kernel size for the first conv layer.
+        dilation_rate (int): dilations rate to increase dilation per layer.
+            If it is 2, dilations are 1, 2, 4, 8 for the next 4 layers.
+        num_blocks (int): number of wavenet blocks.
+        num_layers (int): number of wavenet layers.
+        c_in_channels (int): number of channels of conditioning input.
+        dropout_p (float): dropout rate.
+        weight_norm (bool): enable/disable weight norm for convolution layers.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        num_blocks,
+        num_layers,
+        c_in_channels=0,
+        dropout_p=0,
+        weight_norm=True,
+    ):
+        super().__init__()
+        self.wn_blocks = nn.ModuleList()
+        for idx in range(num_blocks):
+            layer = WN(
+                in_channels=in_channels if idx == 0 else hidden_channels,
+                hidden_channels=hidden_channels,
+                kernel_size=kernel_size,
+                dilation_rate=dilation_rate,
+                num_layers=num_layers,
+                c_in_channels=c_in_channels,
+                dropout_p=dropout_p,
+                weight_norm=weight_norm,
+            )
+            self.wn_blocks.append(layer)
+
+    def forward(self, x, x_mask=None, g=None):
+        o = x
+        for layer in self.wn_blocks:
+            o = layer(o, x_mask, g)
+        return o

TTS/tts/layers/glow_tts/decoder.py

@@ -0,0 +1,141 @@
+import torch
+from torch import nn
+
+from TTS.tts.layers.generic.normalization import ActNorm
+from TTS.tts.layers.glow_tts.glow import CouplingBlock, InvConvNear
+
+
+def squeeze(x, x_mask=None, num_sqz=2):
+    """GlowTTS squeeze operation
+    Increase number of channels and reduce number of time steps
+    by the same factor.
+
+    Note:
+        each 's' is a n-dimensional vector.
+        ``[s1,s2,s3,s4,s5,s6] --> [[s1, s3, s5], [s2, s4, s6]]``
+    """
+    b, c, t = x.size()
+
+    t = (t // num_sqz) * num_sqz
+    x = x[:, :, :t]
+    x_sqz = x.view(b, c, t // num_sqz, num_sqz)
+    x_sqz = x_sqz.permute(0, 3, 1, 2).contiguous().view(b, c * num_sqz, t // num_sqz)
+
+    if x_mask is not None:
+        x_mask = x_mask[:, :, num_sqz - 1 :: num_sqz]
+    else:
+        x_mask = torch.ones(b, 1, t // num_sqz).to(device=x.device, dtype=x.dtype)
+    return x_sqz * x_mask, x_mask
+
+
+def unsqueeze(x, x_mask=None, num_sqz=2):
+    """GlowTTS unsqueeze operation (revert the squeeze)
+
+    Note:
+        each 's' is a n-dimensional vector.
+        ``[[s1, s3, s5], [s2, s4, s6]] --> [[s1, s3, s5, s2, s4, s6]]``
+    """
+    b, c, t = x.size()
+
+    x_unsqz = x.view(b, num_sqz, c // num_sqz, t)
+    x_unsqz = x_unsqz.permute(0, 2, 3, 1).contiguous().view(b, c // num_sqz, t * num_sqz)
+
+    if x_mask is not None:
+        x_mask = x_mask.unsqueeze(-1).repeat(1, 1, 1, num_sqz).view(b, 1, t * num_sqz)
+    else:
+        x_mask = torch.ones(b, 1, t * num_sqz).to(device=x.device, dtype=x.dtype)
+    return x_unsqz * x_mask, x_mask
+
+
+class Decoder(nn.Module):
+    """Stack of Glow Decoder Modules.
+
+    ::
+
+        Squeeze -> ActNorm -> InvertibleConv1x1 -> AffineCoupling -> Unsqueeze
+
+    Args:
+        in_channels (int): channels of input tensor.
+        hidden_channels (int): hidden decoder channels.
+        kernel_size (int): Coupling block kernel size. (Wavenet filter kernel size.)
+        dilation_rate (int): rate to increase dilation by each layer in a decoder block.
+        num_flow_blocks (int): number of decoder blocks.
+        num_coupling_layers (int): number coupling layers. (number of wavenet layers.)
+        dropout_p (float): wavenet dropout rate.
+        sigmoid_scale (bool): enable/disable sigmoid scaling in coupling layer.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        num_flow_blocks,
+        num_coupling_layers,
+        dropout_p=0.0,
+        num_splits=4,
+        num_squeeze=2,
+        sigmoid_scale=False,
+        c_in_channels=0,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.num_flow_blocks = num_flow_blocks
+        self.num_coupling_layers = num_coupling_layers
+        self.dropout_p = dropout_p
+        self.num_splits = num_splits
+        self.num_squeeze = num_squeeze
+        self.sigmoid_scale = sigmoid_scale
+        self.c_in_channels = c_in_channels
+
+        self.flows = nn.ModuleList()
+        for _ in range(num_flow_blocks):
+            self.flows.append(ActNorm(channels=in_channels * num_squeeze))
+            self.flows.append(InvConvNear(channels=in_channels * num_squeeze, num_splits=num_splits))
+            self.flows.append(
+                CouplingBlock(
+                    in_channels * num_squeeze,
+                    hidden_channels,
+                    kernel_size=kernel_size,
+                    dilation_rate=dilation_rate,
+                    num_layers=num_coupling_layers,
+                    c_in_channels=c_in_channels,
+                    dropout_p=dropout_p,
+                    sigmoid_scale=sigmoid_scale,
+                )
+            )
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        """
+        Shapes:
+            - x:  :math:`[B, C, T]`
+            - x_mask: :math:`[B, 1 ,T]`
+            - g: :math:`[B, C]`
+        """
+        if not reverse:
+            flows = self.flows
+            logdet_tot = 0
+        else:
+            flows = reversed(self.flows)
+            logdet_tot = None
+
+        if self.num_squeeze > 1:
+            x, x_mask = squeeze(x, x_mask, self.num_squeeze)
+        for f in flows:
+            if not reverse:
+                x, logdet = f(x, x_mask, g=g, reverse=reverse)
+                logdet_tot += logdet
+            else:
+                x, logdet = f(x, x_mask, g=g, reverse=reverse)
+        if self.num_squeeze > 1:
+            x, x_mask = unsqueeze(x, x_mask, self.num_squeeze)
+        return x, logdet_tot
+
+    def store_inverse(self):
+        for f in self.flows:
+            f.store_inverse()

TTS/tts/layers/glow_tts/duration_predictor.py

@@ -0,0 +1,69 @@
+import torch
+from torch import nn
+
+from ..generic.normalization import LayerNorm
+
+
+class DurationPredictor(nn.Module):
+    """Glow-TTS duration prediction model.
+
+    ::
+
+        [2 x (conv1d_kxk -> relu -> layer_norm -> dropout)] -> conv1d_1x1 -> durs
+
+    Args:
+        in_channels (int): Number of channels of the input tensor.
+        hidden_channels (int): Number of hidden channels of the network.
+        kernel_size (int): Kernel size for the conv layers.
+        dropout_p (float): Dropout rate used after each conv layer.
+    """
+
+    def __init__(self, in_channels, hidden_channels, kernel_size, dropout_p, cond_channels=None, language_emb_dim=None):
+        super().__init__()
+
+        # add language embedding dim in the input
+        if language_emb_dim:
+            in_channels += language_emb_dim
+
+        # class arguments
+        self.in_channels = in_channels
+        self.filter_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dropout_p = dropout_p
+        # layers
+        self.drop = nn.Dropout(dropout_p)
+        self.conv_1 = nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size // 2)
+        self.norm_1 = LayerNorm(hidden_channels)
+        self.conv_2 = nn.Conv1d(hidden_channels, hidden_channels, kernel_size, padding=kernel_size // 2)
+        self.norm_2 = LayerNorm(hidden_channels)
+        # output layer
+        self.proj = nn.Conv1d(hidden_channels, 1, 1)
+        if cond_channels is not None and cond_channels != 0:
+            self.cond = nn.Conv1d(cond_channels, in_channels, 1)
+
+        if language_emb_dim != 0 and language_emb_dim is not None:
+            self.cond_lang = nn.Conv1d(language_emb_dim, in_channels, 1)
+
+    def forward(self, x, x_mask, g=None, lang_emb=None):
+        """
+        Shapes:
+            - x: :math:`[B, C, T]`
+            - x_mask: :math:`[B, 1, T]`
+            - g: :math:`[B, C, 1]`
+        """
+        if g is not None:
+            x = x + self.cond(g)
+
+        if lang_emb is not None:
+            x = x + self.cond_lang(lang_emb)
+
+        x = self.conv_1(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_1(x)
+        x = self.drop(x)
+        x = self.conv_2(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_2(x)
+        x = self.drop(x)
+        x = self.proj(x * x_mask)
+        return x * x_mask

TTS/tts/layers/glow_tts/encoder.py

@@ -0,0 +1,179 @@
+import math
+
+import torch
+from torch import nn
+
+from TTS.tts.layers.generic.gated_conv import GatedConvBlock
+from TTS.tts.layers.generic.res_conv_bn import ResidualConv1dBNBlock
+from TTS.tts.layers.generic.time_depth_sep_conv import TimeDepthSeparableConvBlock
+from TTS.tts.layers.glow_tts.duration_predictor import DurationPredictor
+from TTS.tts.layers.glow_tts.glow import ResidualConv1dLayerNormBlock
+from TTS.tts.layers.glow_tts.transformer import RelativePositionTransformer
+from TTS.tts.utils.helpers import sequence_mask
+
+
+class Encoder(nn.Module):
+    """Glow-TTS encoder module.
+
+    ::
+
+        embedding -> <prenet> -> encoder_module -> <postnet> --> proj_mean
+                                                             |
+                                                             |-> proj_var
+                                                             |
+                                                             |-> concat -> duration_predictor
+                                                                    ↑
+                                                              speaker_embed
+
+    Args:
+        num_chars (int): number of characters.
+        out_channels (int): number of output channels.
+        hidden_channels (int): encoder's embedding size.
+        hidden_channels_ffn (int): transformer's feed-forward channels.
+        kernel_size (int): kernel size for conv layers and duration predictor.
+        dropout_p (float): dropout rate for any dropout layer.
+        mean_only (bool): if True, output only mean values and use constant std.
+        use_prenet (bool): if True, use pre-convolutional layers before transformer layers.
+        c_in_channels (int): number of channels in conditional input.
+
+    Shapes:
+        - input: (B, T, C)
+
+    ::
+
+        suggested encoder params...
+
+        for encoder_type == 'rel_pos_transformer'
+            encoder_params={
+                'kernel_size':3,
+                'dropout_p': 0.1,
+                'num_layers': 6,
+                'num_heads': 2,
+                'hidden_channels_ffn': 768,  # 4 times the hidden_channels
+                'input_length': None
+            }
+
+        for encoder_type == 'gated_conv'
+            encoder_params={
+                'kernel_size':5,
+                'dropout_p': 0.1,
+                'num_layers': 9,
+            }
+
+        for encoder_type == 'residual_conv_bn'
+            encoder_params={
+                "kernel_size": 4,
+                "dilations": [1, 2, 4, 1, 2, 4, 1, 2, 4, 1, 2, 4, 1],
+                "num_conv_blocks": 2,
+                "num_res_blocks": 13
+            }
+
+         for encoder_type == 'time_depth_separable'
+            encoder_params={
+                "kernel_size": 5,
+                'num_layers': 9,
+            }
+    """
+
+    def __init__(
+        self,
+        num_chars,
+        out_channels,
+        hidden_channels,
+        hidden_channels_dp,
+        encoder_type,
+        encoder_params,
+        dropout_p_dp=0.1,
+        mean_only=False,
+        use_prenet=True,
+        c_in_channels=0,
+    ):
+        super().__init__()
+        # class arguments
+        self.num_chars = num_chars
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.hidden_channels_dp = hidden_channels_dp
+        self.dropout_p_dp = dropout_p_dp
+        self.mean_only = mean_only
+        self.use_prenet = use_prenet
+        self.c_in_channels = c_in_channels
+        self.encoder_type = encoder_type
+        # embedding layer
+        self.emb = nn.Embedding(num_chars, hidden_channels)
+        nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)
+        # init encoder module
+        if encoder_type.lower() == "rel_pos_transformer":
+            if use_prenet:
+                self.prenet = ResidualConv1dLayerNormBlock(
+                    hidden_channels, hidden_channels, hidden_channels, kernel_size=5, num_layers=3, dropout_p=0.5
+                )
+            self.encoder = RelativePositionTransformer(
+                hidden_channels, hidden_channels, hidden_channels, **encoder_params
+            )
+        elif encoder_type.lower() == "gated_conv":
+            self.encoder = GatedConvBlock(hidden_channels, **encoder_params)
+        elif encoder_type.lower() == "residual_conv_bn":
+            if use_prenet:
+                self.prenet = nn.Sequential(nn.Conv1d(hidden_channels, hidden_channels, 1), nn.ReLU())
+            self.encoder = ResidualConv1dBNBlock(hidden_channels, hidden_channels, hidden_channels, **encoder_params)
+            self.postnet = nn.Sequential(
+                nn.Conv1d(self.hidden_channels, self.hidden_channels, 1), nn.BatchNorm1d(self.hidden_channels)
+            )
+        elif encoder_type.lower() == "time_depth_separable":
+            if use_prenet:
+                self.prenet = ResidualConv1dLayerNormBlock(
+                    hidden_channels, hidden_channels, hidden_channels, kernel_size=5, num_layers=3, dropout_p=0.5
+                )
+            self.encoder = TimeDepthSeparableConvBlock(
+                hidden_channels, hidden_channels, hidden_channels, **encoder_params
+            )
+        else:
+            raise ValueError(" [!] Unkown encoder type.")
+
+        # final projection layers
+        self.proj_m = nn.Conv1d(hidden_channels, out_channels, 1)
+        if not mean_only:
+            self.proj_s = nn.Conv1d(hidden_channels, out_channels, 1)
+        # duration predictor
+        self.duration_predictor = DurationPredictor(
+            hidden_channels + c_in_channels, hidden_channels_dp, 3, dropout_p_dp
+        )
+
+    def forward(self, x, x_lengths, g=None):
+        """
+        Shapes:
+            - x: :math:`[B, C, T]`
+            - x_lengths: :math:`[B]`
+            - g (optional): :math:`[B, 1, T]`
+        """
+        # embedding layer
+        # [B ,T, D]
+        x = self.emb(x) * math.sqrt(self.hidden_channels)
+        # [B, D, T]
+        x = torch.transpose(x, 1, -1)
+        # compute input sequence mask
+        x_mask = torch.unsqueeze(sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
+        # prenet
+        if hasattr(self, "prenet") and self.use_prenet:
+            x = self.prenet(x, x_mask)
+        # encoder
+        x = self.encoder(x, x_mask)
+        # postnet
+        if hasattr(self, "postnet"):
+            x = self.postnet(x) * x_mask
+        # set duration predictor input
+        if g is not None:
+            g_exp = g.expand(-1, -1, x.size(-1))
+            x_dp = torch.cat([x.detach(), g_exp], 1)
+        else:
+            x_dp = x.detach()
+        # final projection layer
+        x_m = self.proj_m(x) * x_mask
+        if not self.mean_only:
+            x_logs = self.proj_s(x) * x_mask
+        else:
+            x_logs = torch.zeros_like(x_m)
+        # duration predictor
+        logw = self.duration_predictor(x_dp, x_mask)
+        return x_m, x_logs, logw, x_mask

TTS/tts/layers/glow_tts/glow.py

@@ -0,0 +1,233 @@
+import torch
+from packaging.version import Version
+from torch import nn
+from torch.nn import functional as F
+
+from TTS.tts.layers.generic.wavenet import WN
+
+from ..generic.normalization import LayerNorm
+
+
+class ResidualConv1dLayerNormBlock(nn.Module):
+    """Conv1d with Layer Normalization and residual connection as in GlowTTS paper.
+    https://arxiv.org/pdf/1811.00002.pdf
+
+    ::
+
+        x |-> conv1d -> layer_norm -> relu -> dropout -> + -> o
+          |---------------> conv1d_1x1 ------------------|
+
+    Args:
+        in_channels (int): number of input tensor channels.
+        hidden_channels (int): number of inner layer channels.
+        out_channels (int): number of output tensor channels.
+        kernel_size (int): kernel size of conv1d filter.
+        num_layers (int): number of blocks.
+        dropout_p (float): dropout rate for each block.
+    """
+
+    def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, num_layers, dropout_p):
+        super().__init__()
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.num_layers = num_layers
+        self.dropout_p = dropout_p
+        assert num_layers > 1, " [!] number of layers should be > 0."
+        assert kernel_size % 2 == 1, " [!] kernel size should be odd number."
+
+        self.conv_layers = nn.ModuleList()
+        self.norm_layers = nn.ModuleList()
+
+        for idx in range(num_layers):
+            self.conv_layers.append(
+                nn.Conv1d(
+                    in_channels if idx == 0 else hidden_channels, hidden_channels, kernel_size, padding=kernel_size // 2
+                )
+            )
+            self.norm_layers.append(LayerNorm(hidden_channels))
+
+        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
+
+    def forward(self, x, x_mask):
+        """
+        Shapes:
+            - x: :math:`[B, C, T]`
+            - x_mask: :math:`[B, 1, T]`
+        """
+        x_res = x
+        for i in range(self.num_layers):
+            x = self.conv_layers[i](x * x_mask)
+            x = self.norm_layers[i](x * x_mask)
+            x = F.dropout(F.relu(x), self.dropout_p, training=self.training)
+        x = x_res + self.proj(x)
+        return x * x_mask
+
+
+class InvConvNear(nn.Module):
+    """Invertible Convolution with input splitting as in GlowTTS paper.
+    https://arxiv.org/pdf/1811.00002.pdf
+
+    Args:
+        channels (int): input and output channels.
+        num_splits (int): number of splits, also H and W of conv layer.
+        no_jacobian (bool): enable/disable jacobian computations.
+
+    Note:
+        Split the input into groups of size self.num_splits and
+        perform 1x1 convolution separately. Cast 1x1 conv operation
+        to 2d by reshaping the input for efficiency.
+    """
+
+    def __init__(self, channels, num_splits=4, no_jacobian=False, **kwargs):  # pylint: disable=unused-argument
+        super().__init__()
+        assert num_splits % 2 == 0
+        self.channels = channels
+        self.num_splits = num_splits
+        self.no_jacobian = no_jacobian
+        self.weight_inv = None
+
+        if Version(torch.__version__) < Version("1.9"):
+            w_init = torch.qr(torch.FloatTensor(self.num_splits, self.num_splits).normal_())[0]
+        else:
+            w_init = torch.linalg.qr(torch.FloatTensor(self.num_splits, self.num_splits).normal_(), "complete")[0]
+
+        if torch.det(w_init) < 0:
+            w_init[:, 0] = -1 * w_init[:, 0]
+        self.weight = nn.Parameter(w_init)
+
+    def forward(self, x, x_mask=None, reverse=False, **kwargs):  # pylint: disable=unused-argument
+        """
+        Shapes:
+            - x: :math:`[B, C, T]`
+            - x_mask: :math:`[B, 1, T]`
+        """
+        b, c, t = x.size()
+        assert c % self.num_splits == 0
+        if x_mask is None:
+            x_mask = 1
+            x_len = torch.ones((b,), dtype=x.dtype, device=x.device) * t
+        else:
+            x_len = torch.sum(x_mask, [1, 2])
+
+        x = x.view(b, 2, c // self.num_splits, self.num_splits // 2, t)
+        x = x.permute(0, 1, 3, 2, 4).contiguous().view(b, self.num_splits, c // self.num_splits, t)
+
+        if reverse:
+            if self.weight_inv is not None:
+                weight = self.weight_inv
+            else:
+                weight = torch.inverse(self.weight.float()).to(dtype=self.weight.dtype)
+            logdet = None
+        else:
+            weight = self.weight
+            if self.no_jacobian:
+                logdet = 0
+            else:
+                logdet = torch.logdet(self.weight) * (c / self.num_splits) * x_len  # [b]
+
+        weight = weight.view(self.num_splits, self.num_splits, 1, 1)
+        z = F.conv2d(x, weight)
+
+        z = z.view(b, 2, self.num_splits // 2, c // self.num_splits, t)
+        z = z.permute(0, 1, 3, 2, 4).contiguous().view(b, c, t) * x_mask
+        return z, logdet
+
+    def store_inverse(self):
+        weight_inv = torch.inverse(self.weight.float()).to(dtype=self.weight.dtype)
+        self.weight_inv = nn.Parameter(weight_inv, requires_grad=False)
+
+
+class CouplingBlock(nn.Module):
+    """Glow Affine Coupling block as in GlowTTS paper.
+    https://arxiv.org/pdf/1811.00002.pdf
+
+    ::
+
+        x --> x0 -> conv1d -> wavenet -> conv1d --> t, s -> concat(s*x1 + t, x0) -> o
+        '-> x1 - - - - - - - - - - - - - - - - - - - - - - - - - ^
+
+    Args:
+         in_channels (int): number of input tensor channels.
+         hidden_channels (int): number of hidden channels.
+         kernel_size (int): WaveNet filter kernel size.
+         dilation_rate (int): rate to increase dilation by each layer in a decoder block.
+         num_layers (int): number of WaveNet layers.
+         c_in_channels (int): number of conditioning input channels.
+         dropout_p (int): wavenet dropout rate.
+         sigmoid_scale (bool): enable/disable sigmoid scaling for output scale.
+
+    Note:
+         It does not use the conditional inputs differently from WaveGlow.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        num_layers,
+        c_in_channels=0,
+        dropout_p=0,
+        sigmoid_scale=False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.num_layers = num_layers
+        self.c_in_channels = c_in_channels
+        self.dropout_p = dropout_p
+        self.sigmoid_scale = sigmoid_scale
+        # input layer
+        start = torch.nn.Conv1d(in_channels // 2, hidden_channels, 1)
+        start = torch.nn.utils.parametrizations.weight_norm(start)
+        self.start = start
+        # output layer
+        # Initializing last layer to 0 makes the affine coupling layers
+        # do nothing at first.  This helps with training stability
+        end = torch.nn.Conv1d(hidden_channels, in_channels, 1)
+        end.weight.data.zero_()
+        end.bias.data.zero_()
+        self.end = end
+        # coupling layers
+        self.wn = WN(hidden_channels, hidden_channels, kernel_size, dilation_rate, num_layers, c_in_channels, dropout_p)
+
+    def forward(self, x, x_mask=None, reverse=False, g=None, **kwargs):  # pylint: disable=unused-argument
+        """
+        Shapes:
+            - x: :math:`[B, C, T]`
+            - x_mask: :math:`[B, 1, T]`
+            - g: :math:`[B, C, 1]`
+        """
+        if x_mask is None:
+            x_mask = 1
+        x_0, x_1 = x[:, : self.in_channels // 2], x[:, self.in_channels // 2 :]
+
+        x = self.start(x_0) * x_mask
+        x = self.wn(x, x_mask, g)
+        out = self.end(x)
+
+        z_0 = x_0
+        t = out[:, : self.in_channels // 2, :]
+        s = out[:, self.in_channels // 2 :, :]
+        if self.sigmoid_scale:
+            s = torch.log(1e-6 + torch.sigmoid(s + 2))
+
+        if reverse:
+            z_1 = (x_1 - t) * torch.exp(-s) * x_mask
+            logdet = None
+        else:
+            z_1 = (t + torch.exp(s) * x_1) * x_mask
+            logdet = torch.sum(s * x_mask, [1, 2])
+
+        z = torch.cat([z_0, z_1], 1)
+        return z, logdet
+
+    def store_inverse(self):
+        self.wn.remove_weight_norm()

TTS/tts/layers/glow_tts/transformer.py

@@ -0,0 +1,432 @@
+import math
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from TTS.tts.layers.generic.normalization import LayerNorm, LayerNorm2
+
+
+class RelativePositionMultiHeadAttention(nn.Module):
+    """Multi-head attention with Relative Positional embedding.
+    https://arxiv.org/pdf/1809.04281.pdf
+
+    It learns positional embeddings for a window of neighbours. For keys and values,
+    it learns different set of embeddings. Key embeddings are agregated with the attention
+    scores and value embeddings are aggregated with the output.
+
+    Note:
+        Example with relative attention window size 2
+
+        - input = [a, b, c, d, e]
+        - rel_attn_embeddings = [e(t-2), e(t-1), e(t+1), e(t+2)]
+
+        So it learns 4 embedding vectors (in total 8) separately for key and value vectors.
+
+        Considering the input c
+
+        - e(t-2) corresponds to c -> a
+        - e(t-2) corresponds to c -> b
+        - e(t-2) corresponds to c -> d
+        - e(t-2) corresponds to c -> e
+
+        These embeddings are shared among different time steps. So input a, b, d and e also uses
+        the same embeddings.
+
+        Embeddings are ignored when the relative window is out of limit for the first and the last
+        n items.
+
+    Args:
+        channels (int): input and inner layer channels.
+        out_channels (int): output channels.
+        num_heads (int): number of attention heads.
+        rel_attn_window_size (int, optional): relation attention window size.
+            If 4, for each time step next and previous 4 time steps are attended.
+            If default, relative encoding is disabled and it is a regular transformer.
+            Defaults to None.
+        heads_share (bool, optional): [description]. Defaults to True.
+        dropout_p (float, optional): dropout rate. Defaults to 0..
+        input_length (int, optional): intput length for positional encoding. Defaults to None.
+        proximal_bias (bool, optional): enable/disable proximal bias as in the paper. Defaults to False.
+        proximal_init (bool, optional): enable/disable poximal init as in the paper.
+            Init key and query layer weights the same. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        channels,
+        out_channels,
+        num_heads,
+        rel_attn_window_size=None,
+        heads_share=True,
+        dropout_p=0.0,
+        input_length=None,
+        proximal_bias=False,
+        proximal_init=False,
+    ):
+        super().__init__()
+        assert channels % num_heads == 0, " [!] channels should be divisible by num_heads."
+        # class attributes
+        self.channels = channels
+        self.out_channels = out_channels
+        self.num_heads = num_heads
+        self.rel_attn_window_size = rel_attn_window_size
+        self.heads_share = heads_share
+        self.input_length = input_length
+        self.proximal_bias = proximal_bias
+        self.dropout_p = dropout_p
+        self.attn = None
+        # query, key, value layers
+        self.k_channels = channels // num_heads
+        self.conv_q = nn.Conv1d(channels, channels, 1)
+        self.conv_k = nn.Conv1d(channels, channels, 1)
+        self.conv_v = nn.Conv1d(channels, channels, 1)
+        # output layers
+        self.conv_o = nn.Conv1d(channels, out_channels, 1)
+        self.dropout = nn.Dropout(dropout_p)
+        # relative positional encoding layers
+        if rel_attn_window_size is not None:
+            n_heads_rel = 1 if heads_share else num_heads
+            rel_stddev = self.k_channels**-0.5
+            emb_rel_k = nn.Parameter(
+                torch.randn(n_heads_rel, rel_attn_window_size * 2 + 1, self.k_channels) * rel_stddev
+            )
+            emb_rel_v = nn.Parameter(
+                torch.randn(n_heads_rel, rel_attn_window_size * 2 + 1, self.k_channels) * rel_stddev
+            )
+            self.register_parameter("emb_rel_k", emb_rel_k)
+            self.register_parameter("emb_rel_v", emb_rel_v)
+
+        # init layers
+        nn.init.xavier_uniform_(self.conv_q.weight)
+        nn.init.xavier_uniform_(self.conv_k.weight)
+        # proximal bias
+        if proximal_init:
+            self.conv_k.weight.data.copy_(self.conv_q.weight.data)
+            self.conv_k.bias.data.copy_(self.conv_q.bias.data)
+        nn.init.xavier_uniform_(self.conv_v.weight)
+
+    def forward(self, x, c, attn_mask=None):
+        """
+        Shapes:
+            - x: :math:`[B, C, T]`
+            - c: :math:`[B, C, T]`
+            - attn_mask: :math:`[B, 1, T, T]`
+        """
+        q = self.conv_q(x)
+        k = self.conv_k(c)
+        v = self.conv_v(c)
+        x, self.attn = self.attention(q, k, v, mask=attn_mask)
+        x = self.conv_o(x)
+        return x
+
+    def attention(self, query, key, value, mask=None):
+        # reshape [b, d, t] -> [b, n_h, t, d_k]
+        b, d, t_s, t_t = (*key.size(), query.size(2))
+        query = query.view(b, self.num_heads, self.k_channels, t_t).transpose(2, 3)
+        key = key.view(b, self.num_heads, self.k_channels, t_s).transpose(2, 3)
+        value = value.view(b, self.num_heads, self.k_channels, t_s).transpose(2, 3)
+        # compute raw attention scores
+        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.k_channels)
+        # relative positional encoding for scores
+        if self.rel_attn_window_size is not None:
+            assert t_s == t_t, "Relative attention is only available for self-attention."
+            # get relative key embeddings
+            key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
+            rel_logits = self._matmul_with_relative_keys(query, key_relative_embeddings)
+            rel_logits = self._relative_position_to_absolute_position(rel_logits)
+            scores_local = rel_logits / math.sqrt(self.k_channels)
+            scores = scores + scores_local
+        # proximan bias
+        if self.proximal_bias:
+            assert t_s == t_t, "Proximal bias is only available for self-attention."
+            scores = scores + self._attn_proximity_bias(t_s).to(device=scores.device, dtype=scores.dtype)
+        # attention score masking
+        if mask is not None:
+            # add small value to prevent oor error.
+            scores = scores.masked_fill(mask == 0, -1e4)
+            if self.input_length is not None:
+                block_mask = torch.ones_like(scores).triu(-1 * self.input_length).tril(self.input_length)
+                scores = scores * block_mask + -1e4 * (1 - block_mask)
+        # attention score normalization
+        p_attn = F.softmax(scores, dim=-1)  # [b, n_h, t_t, t_s]
+        # apply dropout to attention weights
+        p_attn = self.dropout(p_attn)
+        # compute output
+        output = torch.matmul(p_attn, value)
+        # relative positional encoding for values
+        if self.rel_attn_window_size is not None:
+            relative_weights = self._absolute_position_to_relative_position(p_attn)
+            value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, t_s)
+            output = output + self._matmul_with_relative_values(relative_weights, value_relative_embeddings)
+        output = output.transpose(2, 3).contiguous().view(b, d, t_t)  # [b, n_h, t_t, d_k] -> [b, d, t_t]
+        return output, p_attn
+
+    @staticmethod
+    def _matmul_with_relative_values(p_attn, re):
+        """
+        Args:
+            p_attn (Tensor): attention weights.
+            re (Tensor): relative value embedding vector. (a_(i,j)^V)
+
+        Shapes:
+            -p_attn: :math:`[B, H, T, V]`
+            -re: :math:`[H or 1, V, D]`
+            -logits: :math:`[B, H, T, D]`
+        """
+        logits = torch.matmul(p_attn, re.unsqueeze(0))
+        return logits
+
+    @staticmethod
+    def _matmul_with_relative_keys(query, re):
+        """
+        Args:
+            query (Tensor): batch of query vectors. (x*W^Q)
+            re (Tensor): relative key embedding vector. (a_(i,j)^K)
+
+        Shapes:
+            - query: :math:`[B, H, T, D]`
+            - re: :math:`[H or 1, V, D]`
+            - logits: :math:`[B, H, T, V]`
+        """
+        # logits = torch.einsum('bhld, kmd -> bhlm', [query, re.to(query.dtype)])
+        logits = torch.matmul(query, re.unsqueeze(0).transpose(-2, -1))
+        return logits
+
+    def _get_relative_embeddings(self, relative_embeddings, length):
+        """Convert embedding vestors to a tensor of embeddings"""
+        # Pad first before slice to avoid using cond ops.
+        pad_length = max(length - (self.rel_attn_window_size + 1), 0)
+        slice_start_position = max((self.rel_attn_window_size + 1) - length, 0)
+        slice_end_position = slice_start_position + 2 * length - 1
+        if pad_length > 0:
+            padded_relative_embeddings = F.pad(relative_embeddings, [0, 0, pad_length, pad_length, 0, 0])
+        else:
+            padded_relative_embeddings = relative_embeddings
+        used_relative_embeddings = padded_relative_embeddings[:, slice_start_position:slice_end_position]
+        return used_relative_embeddings
+
+    @staticmethod
+    def _relative_position_to_absolute_position(x):
+        """Converts tensor from relative to absolute indexing for local attention.
+        Shapes:
+            x: :math:`[B, C, T, 2 * T - 1]`
+        Returns:
+            A Tensor of shape :math:`[B, C, T, T]`
+        """
+        batch, heads, length, _ = x.size()
+        # Pad to shift from relative to absolute indexing.
+        x = F.pad(x, [0, 1, 0, 0, 0, 0, 0, 0])
+        # Pad extra elements so to add up to shape (len+1, 2*len-1).
+        x_flat = x.view([batch, heads, length * 2 * length])
+        x_flat = F.pad(x_flat, [0, length - 1, 0, 0, 0, 0])
+        # Reshape and slice out the padded elements.
+        x_final = x_flat.view([batch, heads, length + 1, 2 * length - 1])[:, :, :length, length - 1 :]
+        return x_final
+
+    @staticmethod
+    def _absolute_position_to_relative_position(x):
+        """
+        Shapes:
+            - x: :math:`[B, C, T, T]`
+            - ret: :math:`[B, C, T, 2*T-1]`
+        """
+        batch, heads, length, _ = x.size()
+        # padd along column
+        x = F.pad(x, [0, length - 1, 0, 0, 0, 0, 0, 0])
+        x_flat = x.view([batch, heads, length**2 + length * (length - 1)])
+        # add 0's in the beginning that will skew the elements after reshape
+        x_flat = F.pad(x_flat, [length, 0, 0, 0, 0, 0])
+        x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:]
+        return x_final
+
+    @staticmethod
+    def _attn_proximity_bias(length):
+        """Produce an attention mask that discourages distant
+        attention values.
+        Args:
+            length (int): an integer scalar.
+        Returns:
+            a Tensor with shape :math:`[1, 1, T, T]`
+        """
+        # L
+        r = torch.arange(length, dtype=torch.float32)
+        # L x L
+        diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
+        # scale mask values
+        diff = -torch.log1p(torch.abs(diff))
+        # 1 x 1 x L x L
+        return diff.unsqueeze(0).unsqueeze(0)
+
+
+class FeedForwardNetwork(nn.Module):
+    """Feed Forward Inner layers for Transformer.
+
+    Args:
+        in_channels (int): input tensor channels.
+        out_channels (int): output tensor channels.
+        hidden_channels (int): inner layers hidden channels.
+        kernel_size (int): conv1d filter kernel size.
+        dropout_p (float, optional): dropout rate. Defaults to 0.
+    """
+
+    def __init__(self, in_channels, out_channels, hidden_channels, kernel_size, dropout_p=0.0, causal=False):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dropout_p = dropout_p
+
+        if causal:
+            self.padding = self._causal_padding
+        else:
+            self.padding = self._same_padding
+
+        self.conv_1 = nn.Conv1d(in_channels, hidden_channels, kernel_size)
+        self.conv_2 = nn.Conv1d(hidden_channels, out_channels, kernel_size)
+        self.dropout = nn.Dropout(dropout_p)
+
+    def forward(self, x, x_mask):
+        x = self.conv_1(self.padding(x * x_mask))
+        x = torch.relu(x)
+        x = self.dropout(x)
+        x = self.conv_2(self.padding(x * x_mask))
+        return x * x_mask
+
+    def _causal_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l = self.kernel_size - 1
+        pad_r = 0
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, self._pad_shape(padding))
+        return x
+
+    def _same_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l = (self.kernel_size - 1) // 2
+        pad_r = self.kernel_size // 2
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, self._pad_shape(padding))
+        return x
+
+    @staticmethod
+    def _pad_shape(padding):
+        l = padding[::-1]
+        pad_shape = [item for sublist in l for item in sublist]
+        return pad_shape
+
+
+class RelativePositionTransformer(nn.Module):
+    """Transformer with Relative Potional Encoding.
+    https://arxiv.org/abs/1803.02155
+
+    Args:
+        in_channels (int): number of channels of the input tensor.
+        out_chanels (int): number of channels of the output tensor.
+        hidden_channels (int): model hidden channels.
+        hidden_channels_ffn (int): hidden channels of FeedForwardNetwork.
+        num_heads (int): number of attention heads.
+        num_layers (int): number of transformer layers.
+        kernel_size (int, optional): kernel size of feed-forward inner layers. Defaults to 1.
+        dropout_p (float, optional): dropout rate for self-attention and feed-forward inner layers_per_stack. Defaults to 0.
+        rel_attn_window_size (int, optional): relation attention window size.
+            If 4, for each time step next and previous 4 time steps are attended.
+            If default, relative encoding is disabled and it is a regular transformer.
+            Defaults to None.
+        input_length (int, optional): input lenght to limit position encoding. Defaults to None.
+        layer_norm_type (str, optional): type "1" uses torch tensor operations and type "2" uses torch layer_norm
+            primitive. Use type "2", type "1: is for backward compat. Defaults to "1".
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        hidden_channels: int,
+        hidden_channels_ffn: int,
+        num_heads: int,
+        num_layers: int,
+        kernel_size=1,
+        dropout_p=0.0,
+        rel_attn_window_size: int = None,
+        input_length: int = None,
+        layer_norm_type: str = "1",
+    ):
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.hidden_channels_ffn = hidden_channels_ffn
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.kernel_size = kernel_size
+        self.dropout_p = dropout_p
+        self.rel_attn_window_size = rel_attn_window_size
+
+        self.dropout = nn.Dropout(dropout_p)
+        self.attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+
+        for idx in range(self.num_layers):
+            self.attn_layers.append(
+                RelativePositionMultiHeadAttention(
+                    hidden_channels if idx != 0 else in_channels,
+                    hidden_channels,
+                    num_heads,
+                    rel_attn_window_size=rel_attn_window_size,
+                    dropout_p=dropout_p,
+                    input_length=input_length,
+                )
+            )
+            if layer_norm_type == "1":
+                self.norm_layers_1.append(LayerNorm(hidden_channels))
+            elif layer_norm_type == "2":
+                self.norm_layers_1.append(LayerNorm2(hidden_channels))
+            else:
+                raise ValueError(" [!] Unknown layer norm type")
+
+            if hidden_channels != out_channels and (idx + 1) == self.num_layers:
+                self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
+
+            self.ffn_layers.append(
+                FeedForwardNetwork(
+                    hidden_channels,
+                    hidden_channels if (idx + 1) != self.num_layers else out_channels,
+                    hidden_channels_ffn,
+                    kernel_size,
+                    dropout_p=dropout_p,
+                )
+            )
+
+            if layer_norm_type == "1":
+                self.norm_layers_2.append(LayerNorm(hidden_channels if (idx + 1) != self.num_layers else out_channels))
+            elif layer_norm_type == "2":
+                self.norm_layers_2.append(LayerNorm2(hidden_channels if (idx + 1) != self.num_layers else out_channels))
+            else:
+                raise ValueError(" [!] Unknown layer norm type")
+
+    def forward(self, x, x_mask):
+        """
+        Shapes:
+            - x: :math:`[B, C, T]`
+            - x_mask: :math:`[B, 1, T]`
+        """
+        attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        for i in range(self.num_layers):
+            x = x * x_mask
+            y = self.attn_layers[i](x, x, attn_mask)
+            y = self.dropout(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.dropout(y)
+
+            if (i + 1) == self.num_layers and hasattr(self, "proj"):
+                x = self.proj(x)
+
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x

TTS/tts/layers/losses.py

@@ -0,0 +1,889 @@
+import math
+
+import numpy as np
+import torch
+from coqpit import Coqpit
+from torch import nn
+from torch.nn import functional
+
+from TTS.tts.utils.helpers import sequence_mask
+from TTS.tts.utils.ssim import SSIMLoss as _SSIMLoss
+from TTS.utils.audio.torch_transforms import TorchSTFT
+
+
+# pylint: disable=abstract-method
+# relates https://github.com/pytorch/pytorch/issues/42305
+class L1LossMasked(nn.Module):
+    def __init__(self, seq_len_norm):
+        super().__init__()
+        self.seq_len_norm = seq_len_norm
+
+    def forward(self, x, target, length):
+        """
+        Args:
+            x: A Variable containing a FloatTensor of size
+                (batch, max_len, dim) which contains the
+                unnormalized probability for each class.
+            target: A Variable containing a LongTensor of size
+                (batch, max_len, dim) which contains the index of the true
+                class for each corresponding step.
+            length: A Variable containing a LongTensor of size (batch,)
+                which contains the length of each data in a batch.
+        Shapes:
+            x: B x T X D
+            target: B x T x D
+            length: B
+        Returns:
+            loss: An average loss value in range [0, 1] masked by the length.
+        """
+        # mask: (batch, max_len, 1)
+        target.requires_grad = False
+        mask = sequence_mask(sequence_length=length, max_len=target.size(1)).unsqueeze(2).float()
+        if self.seq_len_norm:
+            norm_w = mask / mask.sum(dim=1, keepdim=True)
+            out_weights = norm_w.div(target.shape[0] * target.shape[2])
+            mask = mask.expand_as(x)
+            loss = functional.l1_loss(x * mask, target * mask, reduction="none")
+            loss = loss.mul(out_weights.to(loss.device)).sum()
+        else:
+            mask = mask.expand_as(x)
+            loss = functional.l1_loss(x * mask, target * mask, reduction="sum")
+            loss = loss / mask.sum()
+        return loss
+
+
+class MSELossMasked(nn.Module):
+    def __init__(self, seq_len_norm):
+        super().__init__()
+        self.seq_len_norm = seq_len_norm
+
+    def forward(self, x, target, length):
+        """
+        Args:
+            x: A Variable containing a FloatTensor of size
+                (batch, max_len, dim) which contains the
+                unnormalized probability for each class.
+            target: A Variable containing a LongTensor of size
+                (batch, max_len, dim) which contains the index of the true
+                class for each corresponding step.
+            length: A Variable containing a LongTensor of size (batch,)
+                which contains the length of each data in a batch.
+        Shapes:
+            - x: :math:`[B, T, D]`
+            - target: :math:`[B, T, D]`
+            - length: :math:`B`
+        Returns:
+            loss: An average loss value in range [0, 1] masked by the length.
+        """
+        # mask: (batch, max_len, 1)
+        target.requires_grad = False
+        mask = sequence_mask(sequence_length=length, max_len=target.size(1)).unsqueeze(2).float()
+        if self.seq_len_norm:
+            norm_w = mask / mask.sum(dim=1, keepdim=True)
+            out_weights = norm_w.div(target.shape[0] * target.shape[2])
+            mask = mask.expand_as(x)
+            loss = functional.mse_loss(x * mask, target * mask, reduction="none")
+            loss = loss.mul(out_weights.to(loss.device)).sum()
+        else:
+            mask = mask.expand_as(x)
+            loss = functional.mse_loss(x * mask, target * mask, reduction="sum")
+            loss = loss / mask.sum()
+        return loss
+
+
+def sample_wise_min_max(x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    """Min-Max normalize tensor through first dimension
+    Shapes:
+        - x: :math:`[B, D1, D2]`
+        - m: :math:`[B, D1, 1]`
+    """
+    maximum = torch.amax(x.masked_fill(~mask, 0), dim=(1, 2), keepdim=True)
+    minimum = torch.amin(x.masked_fill(~mask, np.inf), dim=(1, 2), keepdim=True)
+    return (x - minimum) / (maximum - minimum + 1e-8)
+
+
+class SSIMLoss(torch.nn.Module):
+    """SSIM loss as (1 - SSIM)
+    SSIM is explained here https://en.wikipedia.org/wiki/Structural_similarity
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.loss_func = _SSIMLoss()
+
+    def forward(self, y_hat, y, length):
+        """
+        Args:
+            y_hat (tensor): model prediction values.
+            y (tensor): target values.
+            length (tensor): length of each sample in a batch for masking.
+
+        Shapes:
+            y_hat: B x T X D
+            y: B x T x D
+            length: B
+
+         Returns:
+            loss: An average loss value in range [0, 1] masked by the length.
+        """
+        mask = sequence_mask(sequence_length=length, max_len=y.size(1)).unsqueeze(2)
+        y_norm = sample_wise_min_max(y, mask)
+        y_hat_norm = sample_wise_min_max(y_hat, mask)
+        ssim_loss = self.loss_func((y_norm * mask).unsqueeze(1), (y_hat_norm * mask).unsqueeze(1))
+
+        if ssim_loss.item() > 1.0:
+            print(f" > SSIM loss is out-of-range {ssim_loss.item()}, setting it 1.0")
+            ssim_loss = torch.tensor(1.0, device=ssim_loss.device)
+
+        if ssim_loss.item() < 0.0:
+            print(f" > SSIM loss is out-of-range {ssim_loss.item()}, setting it 0.0")
+            ssim_loss = torch.tensor(0.0, device=ssim_loss.device)
+
+        return ssim_loss
+
+
+class AttentionEntropyLoss(nn.Module):
+    # pylint: disable=R0201
+    def forward(self, align):
+        """
+        Forces attention to be more decisive by penalizing
+        soft attention weights
+        """
+        entropy = torch.distributions.Categorical(probs=align).entropy()
+        loss = (entropy / np.log(align.shape[1])).mean()
+        return loss
+
+
+class BCELossMasked(nn.Module):
+    """BCE loss with masking.
+
+    Used mainly for stopnet in autoregressive models.
+
+    Args:
+        pos_weight (float): weight for positive samples. If set < 1, penalize early stopping. Defaults to None.
+    """
+
+    def __init__(self, pos_weight: float = None):
+        super().__init__()
+        self.register_buffer("pos_weight", torch.tensor([pos_weight]))
+
+    def forward(self, x, target, length):
+        """
+        Args:
+            x: A Variable containing a FloatTensor of size
+                (batch, max_len) which contains the
+                unnormalized probability for each class.
+            target: A Variable containing a LongTensor of size
+                (batch, max_len) which contains the index of the true
+                class for each corresponding step.
+            length: A Variable containing a LongTensor of size (batch,)
+                which contains the length of each data in a batch.
+        Shapes:
+            x: B x T
+            target: B x T
+            length: B
+        Returns:
+            loss: An average loss value in range [0, 1] masked by the length.
+        """
+        target.requires_grad = False
+        if length is not None:
+            # mask: (batch, max_len, 1)
+            mask = sequence_mask(sequence_length=length, max_len=target.size(1))
+            num_items = mask.sum()
+            loss = functional.binary_cross_entropy_with_logits(
+                x.masked_select(mask),
+                target.masked_select(mask),
+                pos_weight=self.pos_weight.to(x.device),
+                reduction="sum",
+            )
+        else:
+            loss = functional.binary_cross_entropy_with_logits(
+                x, target, pos_weight=self.pos_weight.to(x.device), reduction="sum"
+            )
+            num_items = torch.numel(x)
+        loss = loss / num_items
+        return loss
+
+
+class DifferentialSpectralLoss(nn.Module):
+    """Differential Spectral Loss
+    https://arxiv.org/ftp/arxiv/papers/1909/1909.10302.pdf"""
+
+    def __init__(self, loss_func):
+        super().__init__()
+        self.loss_func = loss_func
+
+    def forward(self, x, target, length=None):
+        """
+         Shapes:
+            x: B x T
+            target: B x T
+            length: B
+        Returns:
+            loss: An average loss value in range [0, 1] masked by the length.
+        """
+        x_diff = x[:, 1:] - x[:, :-1]
+        target_diff = target[:, 1:] - target[:, :-1]
+        if length is None:
+            return self.loss_func(x_diff, target_diff)
+        return self.loss_func(x_diff, target_diff, length - 1)
+
+
+class GuidedAttentionLoss(torch.nn.Module):
+    def __init__(self, sigma=0.4):
+        super().__init__()
+        self.sigma = sigma
+
+    def _make_ga_masks(self, ilens, olens):
+        B = len(ilens)
+        max_ilen = max(ilens)
+        max_olen = max(olens)
+        ga_masks = torch.zeros((B, max_olen, max_ilen))
+        for idx, (ilen, olen) in enumerate(zip(ilens, olens)):
+            ga_masks[idx, :olen, :ilen] = self._make_ga_mask(ilen, olen, self.sigma)
+        return ga_masks
+
+    def forward(self, att_ws, ilens, olens):
+        ga_masks = self._make_ga_masks(ilens, olens).to(att_ws.device)
+        seq_masks = self._make_masks(ilens, olens).to(att_ws.device)
+        losses = ga_masks * att_ws
+        loss = torch.mean(losses.masked_select(seq_masks))
+        return loss
+
+    @staticmethod
+    def _make_ga_mask(ilen, olen, sigma):
+        grid_x, grid_y = torch.meshgrid(torch.arange(olen).to(olen), torch.arange(ilen).to(ilen))
+        grid_x, grid_y = grid_x.float(), grid_y.float()
+        return 1.0 - torch.exp(-((grid_y / ilen - grid_x / olen) ** 2) / (2 * (sigma**2)))
+
+    @staticmethod
+    def _make_masks(ilens, olens):
+        in_masks = sequence_mask(ilens)
+        out_masks = sequence_mask(olens)
+        return out_masks.unsqueeze(-1) & in_masks.unsqueeze(-2)
+
+
+class Huber(nn.Module):
+    # pylint: disable=R0201
+    def forward(self, x, y, length=None):
+        """
+        Shapes:
+            x: B x T
+            y: B x T
+            length: B
+        """
+        mask = sequence_mask(sequence_length=length, max_len=y.size(1)).unsqueeze(2).float()
+        return torch.nn.functional.smooth_l1_loss(x * mask, y * mask, reduction="sum") / mask.sum()
+
+
+class ForwardSumLoss(nn.Module):
+    def __init__(self, blank_logprob=-1):
+        super().__init__()
+        self.log_softmax = torch.nn.LogSoftmax(dim=3)
+        self.ctc_loss = torch.nn.CTCLoss(zero_infinity=True)
+        self.blank_logprob = blank_logprob
+
+    def forward(self, attn_logprob, in_lens, out_lens):
+        key_lens = in_lens
+        query_lens = out_lens
+        attn_logprob_padded = torch.nn.functional.pad(input=attn_logprob, pad=(1, 0), value=self.blank_logprob)
+
+        total_loss = 0.0
+        for bid in range(attn_logprob.shape[0]):
+            target_seq = torch.arange(1, key_lens[bid] + 1).unsqueeze(0)
+            curr_logprob = attn_logprob_padded[bid].permute(1, 0, 2)[: query_lens[bid], :, : key_lens[bid] + 1]
+
+            curr_logprob = self.log_softmax(curr_logprob[None])[0]
+            loss = self.ctc_loss(
+                curr_logprob,
+                target_seq,
+                input_lengths=query_lens[bid : bid + 1],
+                target_lengths=key_lens[bid : bid + 1],
+            )
+            total_loss = total_loss + loss
+
+        total_loss = total_loss / attn_logprob.shape[0]
+        return total_loss
+
+
+########################
+# MODEL LOSS LAYERS
+########################
+
+
+class TacotronLoss(torch.nn.Module):
+    """Collection of Tacotron set-up based on provided config."""
+
+    def __init__(self, c, ga_sigma=0.4):
+        super().__init__()
+        self.stopnet_pos_weight = c.stopnet_pos_weight
+        self.use_capacitron_vae = c.use_capacitron_vae
+        if self.use_capacitron_vae:
+            self.capacitron_capacity = c.capacitron_vae.capacitron_capacity
+            self.capacitron_vae_loss_alpha = c.capacitron_vae.capacitron_VAE_loss_alpha
+        self.ga_alpha = c.ga_alpha
+        self.decoder_diff_spec_alpha = c.decoder_diff_spec_alpha
+        self.postnet_diff_spec_alpha = c.postnet_diff_spec_alpha
+        self.decoder_alpha = c.decoder_loss_alpha
+        self.postnet_alpha = c.postnet_loss_alpha
+        self.decoder_ssim_alpha = c.decoder_ssim_alpha
+        self.postnet_ssim_alpha = c.postnet_ssim_alpha
+        self.config = c
+
+        # postnet and decoder loss
+        if c.loss_masking:
+            self.criterion = L1LossMasked(c.seq_len_norm) if c.model in ["Tacotron"] else MSELossMasked(c.seq_len_norm)
+        else:
+            self.criterion = nn.L1Loss() if c.model in ["Tacotron"] else nn.MSELoss()
+        # guided attention loss
+        if c.ga_alpha > 0:
+            self.criterion_ga = GuidedAttentionLoss(sigma=ga_sigma)
+        # differential spectral loss
+        if c.postnet_diff_spec_alpha > 0 or c.decoder_diff_spec_alpha > 0:
+            self.criterion_diff_spec = DifferentialSpectralLoss(loss_func=self.criterion)
+        # ssim loss
+        if c.postnet_ssim_alpha > 0 or c.decoder_ssim_alpha > 0:
+            self.criterion_ssim = SSIMLoss()
+        # stopnet loss
+        # pylint: disable=not-callable
+        self.criterion_st = BCELossMasked(pos_weight=torch.tensor(self.stopnet_pos_weight)) if c.stopnet else None
+
+        # For dev pruposes only
+        self.criterion_capacitron_reconstruction_loss = nn.L1Loss(reduction="sum")
+
+    def forward(
+        self,
+        postnet_output,
+        decoder_output,
+        mel_input,
+        linear_input,
+        stopnet_output,
+        stopnet_target,
+        stop_target_length,
+        capacitron_vae_outputs,
+        output_lens,
+        decoder_b_output,
+        alignments,
+        alignment_lens,
+        alignments_backwards,
+        input_lens,
+    ):
+        # decoder outputs linear or mel spectrograms for Tacotron and Tacotron2
+        # the target should be set acccordingly
+        postnet_target = linear_input if self.config.model.lower() in ["tacotron"] else mel_input
+
+        return_dict = {}
+        # remove lengths if no masking is applied
+        if not self.config.loss_masking:
+            output_lens = None
+        # decoder and postnet losses
+        if self.config.loss_masking:
+            if self.decoder_alpha > 0:
+                decoder_loss = self.criterion(decoder_output, mel_input, output_lens)
+            if self.postnet_alpha > 0:
+                postnet_loss = self.criterion(postnet_output, postnet_target, output_lens)
+        else:
+            if self.decoder_alpha > 0:
+                decoder_loss = self.criterion(decoder_output, mel_input)
+            if self.postnet_alpha > 0:
+                postnet_loss = self.criterion(postnet_output, postnet_target)
+        loss = self.decoder_alpha * decoder_loss + self.postnet_alpha * postnet_loss
+        return_dict["decoder_loss"] = decoder_loss
+        return_dict["postnet_loss"] = postnet_loss
+
+        if self.use_capacitron_vae:
+            # extract capacitron vae infos
+            posterior_distribution, prior_distribution, beta = capacitron_vae_outputs
+
+            # KL divergence term between the posterior and the prior
+            kl_term = torch.mean(torch.distributions.kl_divergence(posterior_distribution, prior_distribution))
+
+            # Limit the mutual information between the data and latent space by the variational capacity limit
+            kl_capacity = kl_term - self.capacitron_capacity
+
+            # pass beta through softplus to keep it positive
+            beta = torch.nn.functional.softplus(beta)[0]
+
+            # This is the term going to the main ADAM optimiser, we detach beta because
+            # beta is optimised by a separate, SGD optimiser below
+            capacitron_vae_loss = beta.detach() * kl_capacity
+
+            # normalize the capacitron_vae_loss as in L1Loss or MSELoss.
+            # After this, both the standard loss and capacitron_vae_loss will be in the same scale.
+            # For this reason we don't need use L1Loss and MSELoss in "sum" reduction mode.
+            # Note: the batch is not considered because the L1Loss was calculated in "sum" mode
+            # divided by the batch size, So not dividing the capacitron_vae_loss by B is legitimate.
+
+            # get B T D dimension from input
+            B, T, D = mel_input.size()
+            # normalize
+            if self.config.loss_masking:
+                # if mask loss get T using the mask
+                T = output_lens.sum() / B
+
+            # Only for dev purposes to be able to compare the reconstruction loss with the values in the
+            # original Capacitron paper
+            return_dict["capaciton_reconstruction_loss"] = (
+                self.criterion_capacitron_reconstruction_loss(decoder_output, mel_input) / decoder_output.size(0)
+            ) + kl_capacity
+
+            capacitron_vae_loss = capacitron_vae_loss / (T * D)
+            capacitron_vae_loss = capacitron_vae_loss * self.capacitron_vae_loss_alpha
+
+            # This is the term to purely optimise beta and to pass into the SGD optimizer
+            beta_loss = torch.negative(beta) * kl_capacity.detach()
+
+            loss += capacitron_vae_loss
+
+            return_dict["capacitron_vae_loss"] = capacitron_vae_loss
+            return_dict["capacitron_vae_beta_loss"] = beta_loss
+            return_dict["capacitron_vae_kl_term"] = kl_term
+            return_dict["capacitron_beta"] = beta
+
+        stop_loss = (
+            self.criterion_st(stopnet_output, stopnet_target, stop_target_length)
+            if self.config.stopnet
+            else torch.zeros(1)
+        )
+        loss += stop_loss
+        return_dict["stopnet_loss"] = stop_loss
+
+        # backward decoder loss (if enabled)
+        if self.config.bidirectional_decoder:
+            if self.config.loss_masking:
+                decoder_b_loss = self.criterion(torch.flip(decoder_b_output, dims=(1,)), mel_input, output_lens)
+            else:
+                decoder_b_loss = self.criterion(torch.flip(decoder_b_output, dims=(1,)), mel_input)
+            decoder_c_loss = torch.nn.functional.l1_loss(torch.flip(decoder_b_output, dims=(1,)), decoder_output)
+            loss += self.decoder_alpha * (decoder_b_loss + decoder_c_loss)
+            return_dict["decoder_b_loss"] = decoder_b_loss
+            return_dict["decoder_c_loss"] = decoder_c_loss
+
+        # double decoder consistency loss (if enabled)
+        if self.config.double_decoder_consistency:
+            if self.config.loss_masking:
+                decoder_b_loss = self.criterion(decoder_b_output, mel_input, output_lens)
+            else:
+                decoder_b_loss = self.criterion(decoder_b_output, mel_input)
+            # decoder_c_loss = torch.nn.functional.l1_loss(decoder_b_output, decoder_output)
+            attention_c_loss = torch.nn.functional.l1_loss(alignments, alignments_backwards)
+            loss += self.decoder_alpha * (decoder_b_loss + attention_c_loss)
+            return_dict["decoder_coarse_loss"] = decoder_b_loss
+            return_dict["decoder_ddc_loss"] = attention_c_loss
+
+        # guided attention loss (if enabled)
+        if self.config.ga_alpha > 0:
+            ga_loss = self.criterion_ga(alignments, input_lens, alignment_lens)
+            loss += ga_loss * self.ga_alpha
+            return_dict["ga_loss"] = ga_loss
+
+        # decoder differential spectral loss
+        if self.config.decoder_diff_spec_alpha > 0:
+            decoder_diff_spec_loss = self.criterion_diff_spec(decoder_output, mel_input, output_lens)
+            loss += decoder_diff_spec_loss * self.decoder_diff_spec_alpha
+            return_dict["decoder_diff_spec_loss"] = decoder_diff_spec_loss
+
+        # postnet differential spectral loss
+        if self.config.postnet_diff_spec_alpha > 0:
+            postnet_diff_spec_loss = self.criterion_diff_spec(postnet_output, postnet_target, output_lens)
+            loss += postnet_diff_spec_loss * self.postnet_diff_spec_alpha
+            return_dict["postnet_diff_spec_loss"] = postnet_diff_spec_loss
+
+        # decoder ssim loss
+        if self.config.decoder_ssim_alpha > 0:
+            decoder_ssim_loss = self.criterion_ssim(decoder_output, mel_input, output_lens)
+            loss += decoder_ssim_loss * self.postnet_ssim_alpha
+            return_dict["decoder_ssim_loss"] = decoder_ssim_loss
+
+        # postnet ssim loss
+        if self.config.postnet_ssim_alpha > 0:
+            postnet_ssim_loss = self.criterion_ssim(postnet_output, postnet_target, output_lens)
+            loss += postnet_ssim_loss * self.postnet_ssim_alpha
+            return_dict["postnet_ssim_loss"] = postnet_ssim_loss
+
+        return_dict["loss"] = loss
+        return return_dict
+
+
+class GlowTTSLoss(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.constant_factor = 0.5 * math.log(2 * math.pi)
+
+    def forward(self, z, means, scales, log_det, y_lengths, o_dur_log, o_attn_dur, x_lengths):
+        return_dict = {}
+        # flow loss - neg log likelihood
+        pz = torch.sum(scales) + 0.5 * torch.sum(torch.exp(-2 * scales) * (z - means) ** 2)
+        log_mle = self.constant_factor + (pz - torch.sum(log_det)) / (torch.sum(y_lengths) * z.shape[2])
+        # duration loss - MSE
+        loss_dur = torch.sum((o_dur_log - o_attn_dur) ** 2) / torch.sum(x_lengths)
+        # duration loss - huber loss
+        # loss_dur = torch.nn.functional.smooth_l1_loss(o_dur_log, o_attn_dur, reduction="sum") / torch.sum(x_lengths)
+        return_dict["loss"] = log_mle + loss_dur
+        return_dict["log_mle"] = log_mle
+        return_dict["loss_dur"] = loss_dur
+
+        # check if any loss is NaN
+        for key, loss in return_dict.items():
+            if torch.isnan(loss):
+                raise RuntimeError(f" [!] NaN loss with {key}.")
+        return return_dict
+
+
+def mse_loss_custom(x, y):
+    """MSE loss using the torch back-end without reduction.
+    It uses less VRAM than the raw code"""
+    expanded_x, expanded_y = torch.broadcast_tensors(x, y)
+    return torch._C._nn.mse_loss(expanded_x, expanded_y, 0)  # pylint: disable=protected-access, c-extension-no-member
+
+
+class MDNLoss(nn.Module):
+    """Mixture of Density Network Loss as described in https://arxiv.org/pdf/2003.01950.pdf."""
+
+    def forward(self, logp, text_lengths, mel_lengths):  # pylint: disable=no-self-use
+        """
+        Shapes:
+            mu: [B, D, T]
+            log_sigma: [B, D, T]
+            mel_spec: [B, D, T]
+        """
+        B, T_seq, T_mel = logp.shape
+        log_alpha = logp.new_ones(B, T_seq, T_mel) * (-1e4)
+        log_alpha[:, 0, 0] = logp[:, 0, 0]
+        for t in range(1, T_mel):
+            prev_step = torch.cat(
+                [log_alpha[:, :, t - 1 : t], functional.pad(log_alpha[:, :, t - 1 : t], (0, 0, 1, -1), value=-1e4)],
+                dim=-1,
+            )
+            log_alpha[:, :, t] = torch.logsumexp(prev_step + 1e-4, dim=-1) + logp[:, :, t]
+        alpha_last = log_alpha[torch.arange(B), text_lengths - 1, mel_lengths - 1]
+        mdn_loss = -alpha_last.mean() / T_seq
+        return mdn_loss  # , log_prob_matrix
+
+
+class AlignTTSLoss(nn.Module):
+    """Modified AlignTTS Loss.
+    Computes
+        - L1 and SSIM losses from output spectrograms.
+        - Huber loss for duration predictor.
+        - MDNLoss for Mixture of Density Network.
+
+    All loss values are aggregated by a weighted sum of the alpha values.
+
+    Args:
+        c (dict): TTS model configuration.
+    """
+
+    def __init__(self, c):
+        super().__init__()
+        self.mdn_loss = MDNLoss()
+        self.spec_loss = MSELossMasked(False)
+        self.ssim = SSIMLoss()
+        self.dur_loss = MSELossMasked(False)
+
+        self.ssim_alpha = c.ssim_alpha
+        self.dur_loss_alpha = c.dur_loss_alpha
+        self.spec_loss_alpha = c.spec_loss_alpha
+        self.mdn_alpha = c.mdn_alpha
+
+    def forward(
+        self, logp, decoder_output, decoder_target, decoder_output_lens, dur_output, dur_target, input_lens, phase
+    ):
+        # ssim_alpha, dur_loss_alpha, spec_loss_alpha, mdn_alpha = self.set_alphas(step)
+        spec_loss, ssim_loss, dur_loss, mdn_loss = 0, 0, 0, 0
+        if phase == 0:
+            mdn_loss = self.mdn_loss(logp, input_lens, decoder_output_lens)
+        elif phase == 1:
+            spec_loss = self.spec_loss(decoder_output, decoder_target, decoder_output_lens)
+            ssim_loss = self.ssim(decoder_output, decoder_target, decoder_output_lens)
+        elif phase == 2:
+            mdn_loss = self.mdn_loss(logp, input_lens, decoder_output_lens)
+            spec_loss = self.spec_lossX(decoder_output, decoder_target, decoder_output_lens)
+            ssim_loss = self.ssim(decoder_output, decoder_target, decoder_output_lens)
+        elif phase == 3:
+            dur_loss = self.dur_loss(dur_output.unsqueeze(2), dur_target.unsqueeze(2), input_lens)
+        else:
+            mdn_loss = self.mdn_loss(logp, input_lens, decoder_output_lens)
+            spec_loss = self.spec_loss(decoder_output, decoder_target, decoder_output_lens)
+            ssim_loss = self.ssim(decoder_output, decoder_target, decoder_output_lens)
+            dur_loss = self.dur_loss(dur_output.unsqueeze(2), dur_target.unsqueeze(2), input_lens)
+        loss = (
+            self.spec_loss_alpha * spec_loss
+            + self.ssim_alpha * ssim_loss
+            + self.dur_loss_alpha * dur_loss
+            + self.mdn_alpha * mdn_loss
+        )
+        return {"loss": loss, "loss_l1": spec_loss, "loss_ssim": ssim_loss, "loss_dur": dur_loss, "mdn_loss": mdn_loss}
+
+
+class VitsGeneratorLoss(nn.Module):
+    def __init__(self, c: Coqpit):
+        super().__init__()
+        self.kl_loss_alpha = c.kl_loss_alpha
+        self.gen_loss_alpha = c.gen_loss_alpha
+        self.feat_loss_alpha = c.feat_loss_alpha
+        self.dur_loss_alpha = c.dur_loss_alpha
+        self.mel_loss_alpha = c.mel_loss_alpha
+        self.spk_encoder_loss_alpha = c.speaker_encoder_loss_alpha
+        self.stft = TorchSTFT(
+            c.audio.fft_size,
+            c.audio.hop_length,
+            c.audio.win_length,
+            sample_rate=c.audio.sample_rate,
+            mel_fmin=c.audio.mel_fmin,
+            mel_fmax=c.audio.mel_fmax,
+            n_mels=c.audio.num_mels,
+            use_mel=True,
+            do_amp_to_db=True,
+        )
+
+    @staticmethod
+    def feature_loss(feats_real, feats_generated):
+        loss = 0
+        for dr, dg in zip(feats_real, feats_generated):
+            for rl, gl in zip(dr, dg):
+                rl = rl.float().detach()
+                gl = gl.float()
+                loss += torch.mean(torch.abs(rl - gl))
+        return loss * 2
+
+    @staticmethod
+    def generator_loss(scores_fake):
+        loss = 0
+        gen_losses = []
+        for dg in scores_fake:
+            dg = dg.float()
+            l = torch.mean((1 - dg) ** 2)
+            gen_losses.append(l)
+            loss += l
+
+        return loss, gen_losses
+
+    @staticmethod
+    def kl_loss(z_p, logs_q, m_p, logs_p, z_mask):
+        """
+        z_p, logs_q: [b, h, t_t]
+        m_p, logs_p: [b, h, t_t]
+        """
+        z_p = z_p.float()
+        logs_q = logs_q.float()
+        m_p = m_p.float()
+        logs_p = logs_p.float()
+        z_mask = z_mask.float()
+
+        kl = logs_p - logs_q - 0.5
+        kl += 0.5 * ((z_p - m_p) ** 2) * torch.exp(-2.0 * logs_p)
+        kl = torch.sum(kl * z_mask)
+        l = kl / torch.sum(z_mask)
+        return l
+
+    @staticmethod
+    def cosine_similarity_loss(gt_spk_emb, syn_spk_emb):
+        return -torch.nn.functional.cosine_similarity(gt_spk_emb, syn_spk_emb).mean()
+
+    def forward(
+        self,
+        mel_slice,
+        mel_slice_hat,
+        z_p,
+        logs_q,
+        m_p,
+        logs_p,
+        z_len,
+        scores_disc_fake,
+        feats_disc_fake,
+        feats_disc_real,
+        loss_duration,
+        use_speaker_encoder_as_loss=False,
+        gt_spk_emb=None,
+        syn_spk_emb=None,
+    ):
+        """
+        Shapes:
+            - mel_slice : :math:`[B, 1, T]`
+            - mel_slice_hat: :math:`[B, 1, T]`
+            - z_p: :math:`[B, C, T]`
+            - logs_q: :math:`[B, C, T]`
+            - m_p: :math:`[B, C, T]`
+            - logs_p: :math:`[B, C, T]`
+            - z_len: :math:`[B]`
+            - scores_disc_fake[i]: :math:`[B, C]`
+            - feats_disc_fake[i][j]: :math:`[B, C, T', P]`
+            - feats_disc_real[i][j]: :math:`[B, C, T', P]`
+        """
+        loss = 0.0
+        return_dict = {}
+        z_mask = sequence_mask(z_len).float()
+        # compute losses
+        loss_kl = (
+            self.kl_loss(z_p=z_p, logs_q=logs_q, m_p=m_p, logs_p=logs_p, z_mask=z_mask.unsqueeze(1))
+            * self.kl_loss_alpha
+        )
+        loss_feat = (
+            self.feature_loss(feats_real=feats_disc_real, feats_generated=feats_disc_fake) * self.feat_loss_alpha
+        )
+        loss_gen = self.generator_loss(scores_fake=scores_disc_fake)[0] * self.gen_loss_alpha
+        loss_mel = torch.nn.functional.l1_loss(mel_slice, mel_slice_hat) * self.mel_loss_alpha
+        loss_duration = torch.sum(loss_duration.float()) * self.dur_loss_alpha
+        loss = loss_kl + loss_feat + loss_mel + loss_gen + loss_duration
+
+        if use_speaker_encoder_as_loss:
+            loss_se = self.cosine_similarity_loss(gt_spk_emb, syn_spk_emb) * self.spk_encoder_loss_alpha
+            loss = loss + loss_se
+            return_dict["loss_spk_encoder"] = loss_se
+        # pass losses to the dict
+        return_dict["loss_gen"] = loss_gen
+        return_dict["loss_kl"] = loss_kl
+        return_dict["loss_feat"] = loss_feat
+        return_dict["loss_mel"] = loss_mel
+        return_dict["loss_duration"] = loss_duration
+        return_dict["loss"] = loss
+        return return_dict
+
+
+class VitsDiscriminatorLoss(nn.Module):
+    def __init__(self, c: Coqpit):
+        super().__init__()
+        self.disc_loss_alpha = c.disc_loss_alpha
+
+    @staticmethod
+    def discriminator_loss(scores_real, scores_fake):
+        loss = 0
+        real_losses = []
+        fake_losses = []
+        for dr, dg in zip(scores_real, scores_fake):
+            dr = dr.float()
+            dg = dg.float()
+            real_loss = torch.mean((1 - dr) ** 2)
+            fake_loss = torch.mean(dg**2)
+            loss += real_loss + fake_loss
+            real_losses.append(real_loss.item())
+            fake_losses.append(fake_loss.item())
+        return loss, real_losses, fake_losses
+
+    def forward(self, scores_disc_real, scores_disc_fake):
+        loss = 0.0
+        return_dict = {}
+        loss_disc, loss_disc_real, _ = self.discriminator_loss(
+            scores_real=scores_disc_real, scores_fake=scores_disc_fake
+        )
+        return_dict["loss_disc"] = loss_disc * self.disc_loss_alpha
+        loss = loss + return_dict["loss_disc"]
+        return_dict["loss"] = loss
+
+        for i, ldr in enumerate(loss_disc_real):
+            return_dict[f"loss_disc_real_{i}"] = ldr
+        return return_dict
+
+
+class ForwardTTSLoss(nn.Module):
+    """Generic configurable ForwardTTS loss."""
+
+    def __init__(self, c):
+        super().__init__()
+        if c.spec_loss_type == "mse":
+            self.spec_loss = MSELossMasked(False)
+        elif c.spec_loss_type == "l1":
+            self.spec_loss = L1LossMasked(False)
+        else:
+            raise ValueError(" [!] Unknown spec_loss_type {}".format(c.spec_loss_type))
+
+        if c.duration_loss_type == "mse":
+            self.dur_loss = MSELossMasked(False)
+        elif c.duration_loss_type == "l1":
+            self.dur_loss = L1LossMasked(False)
+        elif c.duration_loss_type == "huber":
+            self.dur_loss = Huber()
+        else:
+            raise ValueError(" [!] Unknown duration_loss_type {}".format(c.duration_loss_type))
+
+        if c.model_args.use_aligner:
+            self.aligner_loss = ForwardSumLoss()
+            self.aligner_loss_alpha = c.aligner_loss_alpha
+
+        if c.model_args.use_pitch:
+            self.pitch_loss = MSELossMasked(False)
+            self.pitch_loss_alpha = c.pitch_loss_alpha
+
+        if c.model_args.use_energy:
+            self.energy_loss = MSELossMasked(False)
+            self.energy_loss_alpha = c.energy_loss_alpha
+
+        if c.use_ssim_loss:
+            self.ssim = SSIMLoss() if c.use_ssim_loss else None
+            self.ssim_loss_alpha = c.ssim_loss_alpha
+
+        self.spec_loss_alpha = c.spec_loss_alpha
+        self.dur_loss_alpha = c.dur_loss_alpha
+        self.binary_alignment_loss_alpha = c.binary_align_loss_alpha
+
+    @staticmethod
+    def _binary_alignment_loss(alignment_hard, alignment_soft):
+        """Binary loss that forces soft alignments to match the hard alignments as
+        explained in `https://arxiv.org/pdf/2108.10447.pdf`.
+        """
+        log_sum = torch.log(torch.clamp(alignment_soft[alignment_hard == 1], min=1e-12)).sum()
+        return -log_sum / alignment_hard.sum()
+
+    def forward(
+        self,
+        decoder_output,
+        decoder_target,
+        decoder_output_lens,
+        dur_output,
+        dur_target,
+        pitch_output,
+        pitch_target,
+        energy_output,
+        energy_target,
+        input_lens,
+        alignment_logprob=None,
+        alignment_hard=None,
+        alignment_soft=None,
+        binary_loss_weight=None,
+    ):
+        loss = 0
+        return_dict = {}
+        if hasattr(self, "ssim_loss") and self.ssim_loss_alpha > 0:
+            ssim_loss = self.ssim(decoder_output, decoder_target, decoder_output_lens)
+            loss = loss + self.ssim_loss_alpha * ssim_loss
+            return_dict["loss_ssim"] = self.ssim_loss_alpha * ssim_loss
+
+        if self.spec_loss_alpha > 0:
+            spec_loss = self.spec_loss(decoder_output, decoder_target, decoder_output_lens)
+            loss = loss + self.spec_loss_alpha * spec_loss
+            return_dict["loss_spec"] = self.spec_loss_alpha * spec_loss
+
+        if self.dur_loss_alpha > 0:
+            log_dur_tgt = torch.log(dur_target.float() + 1)
+            dur_loss = self.dur_loss(dur_output[:, :, None], log_dur_tgt[:, :, None], input_lens)
+            loss = loss + self.dur_loss_alpha * dur_loss
+            return_dict["loss_dur"] = self.dur_loss_alpha * dur_loss
+
+        if hasattr(self, "pitch_loss") and self.pitch_loss_alpha > 0:
+            pitch_loss = self.pitch_loss(pitch_output.transpose(1, 2), pitch_target.transpose(1, 2), input_lens)
+            loss = loss + self.pitch_loss_alpha * pitch_loss
+            return_dict["loss_pitch"] = self.pitch_loss_alpha * pitch_loss
+
+        if hasattr(self, "energy_loss") and self.energy_loss_alpha > 0:
+            energy_loss = self.energy_loss(energy_output.transpose(1, 2), energy_target.transpose(1, 2), input_lens)
+            loss = loss + self.energy_loss_alpha * energy_loss
+            return_dict["loss_energy"] = self.energy_loss_alpha * energy_loss
+
+        if hasattr(self, "aligner_loss") and self.aligner_loss_alpha > 0:
+            aligner_loss = self.aligner_loss(alignment_logprob, input_lens, decoder_output_lens)
+            loss = loss + self.aligner_loss_alpha * aligner_loss
+            return_dict["loss_aligner"] = self.aligner_loss_alpha * aligner_loss
+
+        if self.binary_alignment_loss_alpha > 0 and alignment_hard is not None:
+            binary_alignment_loss = self._binary_alignment_loss(alignment_hard, alignment_soft)
+            loss = loss + self.binary_alignment_loss_alpha * binary_alignment_loss
+            if binary_loss_weight:
+                return_dict["loss_binary_alignment"] = (
+                    self.binary_alignment_loss_alpha * binary_alignment_loss * binary_loss_weight
+                )
+            else:
+                return_dict["loss_binary_alignment"] = self.binary_alignment_loss_alpha * binary_alignment_loss
+
+        return_dict["loss"] = loss
+        return return_dict

TTS/tts/layers/overflow/common_layers.py

@@ -0,0 +1,323 @@
+from typing import List, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from tqdm.auto import tqdm
+
+from TTS.tts.layers.tacotron.common_layers import Linear
+from TTS.tts.layers.tacotron.tacotron2 import ConvBNBlock
+
+
+class Encoder(nn.Module):
+    r"""Neural HMM Encoder
+
+    Same as Tacotron 2 encoder but increases the input length by states per phone
+
+    Args:
+        num_chars (int): Number of characters in the input.
+        state_per_phone (int): Number of states per phone.
+        in_out_channels (int): number of input and output channels.
+        n_convolutions (int): number of convolutional layers.
+    """
+
+    def __init__(self, num_chars, state_per_phone, in_out_channels=512, n_convolutions=3):
+        super().__init__()
+
+        self.state_per_phone = state_per_phone
+        self.in_out_channels = in_out_channels
+
+        self.emb = nn.Embedding(num_chars, in_out_channels)
+        self.convolutions = nn.ModuleList()
+        for _ in range(n_convolutions):
+            self.convolutions.append(ConvBNBlock(in_out_channels, in_out_channels, 5, "relu"))
+        self.lstm = nn.LSTM(
+            in_out_channels,
+            int(in_out_channels / 2) * state_per_phone,
+            num_layers=1,
+            batch_first=True,
+            bias=True,
+            bidirectional=True,
+        )
+        self.rnn_state = None
+
+    def forward(self, x: torch.FloatTensor, x_len: torch.LongTensor) -> Tuple[torch.FloatTensor, torch.LongTensor]:
+        """Forward pass to the encoder.
+
+        Args:
+            x (torch.FloatTensor): input text indices.
+                - shape: :math:`(b, T_{in})`
+            x_len (torch.LongTensor): input text lengths.
+                - shape: :math:`(b,)`
+
+        Returns:
+            Tuple[torch.FloatTensor, torch.LongTensor]: encoder outputs and output lengths.
+                -shape: :math:`((b, T_{in} * states_per_phone, in_out_channels), (b,))`
+        """
+        b, T = x.shape
+        o = self.emb(x).transpose(1, 2)
+        for layer in self.convolutions:
+            o = layer(o)
+        o = o.transpose(1, 2)
+        o = nn.utils.rnn.pack_padded_sequence(o, x_len.cpu(), batch_first=True)
+        self.lstm.flatten_parameters()
+        o, _ = self.lstm(o)
+        o, _ = nn.utils.rnn.pad_packed_sequence(o, batch_first=True)
+        o = o.reshape(b, T * self.state_per_phone, self.in_out_channels)
+        x_len = x_len * self.state_per_phone
+        return o, x_len
+
+    def inference(self, x, x_len):
+        """Inference to the encoder.
+
+        Args:
+            x (torch.FloatTensor): input text indices.
+                - shape: :math:`(b, T_{in})`
+            x_len (torch.LongTensor): input text lengths.
+                - shape: :math:`(b,)`
+
+        Returns:
+            Tuple[torch.FloatTensor, torch.LongTensor]: encoder outputs and output lengths.
+                -shape: :math:`((b, T_{in} * states_per_phone, in_out_channels), (b,))`
+        """
+        b, T = x.shape
+        o = self.emb(x).transpose(1, 2)
+        for layer in self.convolutions:
+            o = layer(o)
+        o = o.transpose(1, 2)
+        # self.lstm.flatten_parameters()
+        o, _ = self.lstm(o)
+        o = o.reshape(b, T * self.state_per_phone, self.in_out_channels)
+        x_len = x_len * self.state_per_phone
+        return o, x_len
+
+
+class ParameterModel(nn.Module):
+    r"""Main neural network of the outputnet
+
+    Note: Do not put dropout layers here, the model will not converge.
+
+    Args:
+            outputnet_size (List[int]): the architecture of the parameter model
+            input_size (int): size of input for the first layer
+            output_size (int): size of output i.e size of the feature dim
+            frame_channels (int): feature dim to set the flat start bias
+            flat_start_params (dict): flat start parameters to set the bias
+    """
+
+    def __init__(
+        self,
+        outputnet_size: List[int],
+        input_size: int,
+        output_size: int,
+        frame_channels: int,
+        flat_start_params: dict,
+    ):
+        super().__init__()
+        self.frame_channels = frame_channels
+
+        self.layers = nn.ModuleList(
+            [Linear(inp, out) for inp, out in zip([input_size] + outputnet_size[:-1], outputnet_size)]
+        )
+        self.last_layer = nn.Linear(outputnet_size[-1], output_size)
+        self.flat_start_output_layer(
+            flat_start_params["mean"], flat_start_params["std"], flat_start_params["transition_p"]
+        )
+
+    def flat_start_output_layer(self, mean, std, transition_p):
+        self.last_layer.weight.data.zero_()
+        self.last_layer.bias.data[0 : self.frame_channels] = mean
+        self.last_layer.bias.data[self.frame_channels : 2 * self.frame_channels] = OverflowUtils.inverse_softplus(std)
+        self.last_layer.bias.data[2 * self.frame_channels :] = OverflowUtils.inverse_sigmod(transition_p)
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = F.relu(layer(x))
+        x = self.last_layer(x)
+        return x
+
+
+class Outputnet(nn.Module):
+    r"""
+    This network takes current state and previous observed values as input
+    and returns its parameters, mean, standard deviation and probability
+    of transition to the next state
+    """
+
+    def __init__(
+        self,
+        encoder_dim: int,
+        memory_rnn_dim: int,
+        frame_channels: int,
+        outputnet_size: List[int],
+        flat_start_params: dict,
+        std_floor: float = 1e-2,
+    ):
+        super().__init__()
+
+        self.frame_channels = frame_channels
+        self.flat_start_params = flat_start_params
+        self.std_floor = std_floor
+
+        input_size = memory_rnn_dim + encoder_dim
+        output_size = 2 * frame_channels + 1
+
+        self.parametermodel = ParameterModel(
+            outputnet_size=outputnet_size,
+            input_size=input_size,
+            output_size=output_size,
+            flat_start_params=flat_start_params,
+            frame_channels=frame_channels,
+        )
+
+    def forward(self, ar_mels, inputs):
+        r"""Inputs observation and returns the means, stds and transition probability for the current state
+
+        Args:
+            ar_mel_inputs (torch.FloatTensor): shape (batch, prenet_dim)
+            states (torch.FloatTensor):  (batch, hidden_states, hidden_state_dim)
+
+        Returns:
+            means: means for the emission observation for each feature
+                - shape: (B, hidden_states, feature_size)
+            stds: standard deviations for the emission observation for each feature
+                - shape: (batch, hidden_states, feature_size)
+            transition_vectors: transition vector for the current hidden state
+                - shape: (batch, hidden_states)
+        """
+        batch_size, prenet_dim = ar_mels.shape[0], ar_mels.shape[1]
+        N = inputs.shape[1]
+
+        ar_mels = ar_mels.unsqueeze(1).expand(batch_size, N, prenet_dim)
+        ar_mels = torch.cat((ar_mels, inputs), dim=2)
+        ar_mels = self.parametermodel(ar_mels)
+
+        mean, std, transition_vector = (
+            ar_mels[:, :, 0 : self.frame_channels],
+            ar_mels[:, :, self.frame_channels : 2 * self.frame_channels],
+            ar_mels[:, :, 2 * self.frame_channels :].squeeze(2),
+        )
+        std = F.softplus(std)
+        std = self._floor_std(std)
+        return mean, std, transition_vector
+
+    def _floor_std(self, std):
+        r"""
+        It clamps the standard deviation to not to go below some level
+        This removes the problem when the model tries to cheat for higher likelihoods by converting
+        one of the gaussians to a point mass.
+
+        Args:
+            std (float Tensor): tensor containing the standard deviation to be
+        """
+        original_tensor = std.clone().detach()
+        std = torch.clamp(std, min=self.std_floor)
+        if torch.any(original_tensor != std):
+            print(
+                "[*] Standard deviation was floored! The model is preventing overfitting, nothing serious to worry about"
+            )
+        return std
+
+
+class OverflowUtils:
+    @staticmethod
+    def get_data_parameters_for_flat_start(
+        data_loader: torch.utils.data.DataLoader, out_channels: int, states_per_phone: int
+    ):
+        """Generates data parameters for flat starting the HMM.
+
+        Args:
+            data_loader (torch.utils.data.Dataloader): _description_
+            out_channels (int): mel spectrogram channels
+            states_per_phone (_type_): HMM states per phone
+        """
+
+        # State related information for transition_p
+        total_state_len = 0
+        total_mel_len = 0
+
+        # Useful for data mean an std
+        total_mel_sum = 0
+        total_mel_sq_sum = 0
+
+        for batch in tqdm(data_loader, leave=False):
+            text_lengths = batch["token_id_lengths"]
+            mels = batch["mel"]
+            mel_lengths = batch["mel_lengths"]
+
+            total_state_len += torch.sum(text_lengths)
+            total_mel_len += torch.sum(mel_lengths)
+            total_mel_sum += torch.sum(mels)
+            total_mel_sq_sum += torch.sum(torch.pow(mels, 2))
+
+        data_mean = total_mel_sum / (total_mel_len * out_channels)
+        data_std = torch.sqrt((total_mel_sq_sum / (total_mel_len * out_channels)) - torch.pow(data_mean, 2))
+        average_num_states = total_state_len / len(data_loader.dataset)
+        average_mel_len = total_mel_len / len(data_loader.dataset)
+        average_duration_each_state = average_mel_len / average_num_states
+        init_transition_prob = 1 / average_duration_each_state
+
+        return data_mean, data_std, (init_transition_prob * states_per_phone)
+
+    @staticmethod
+    @torch.no_grad()
+    def update_flat_start_transition(model, transition_p):
+        model.neural_hmm.output_net.parametermodel.flat_start_output_layer(0.0, 1.0, transition_p)
+
+    @staticmethod
+    def log_clamped(x, eps=1e-04):
+        """
+        Avoids the log(0) problem
+
+        Args:
+            x (torch.tensor): input tensor
+            eps (float, optional): lower bound. Defaults to 1e-04.
+
+        Returns:
+            torch.tensor: :math:`log(x)`
+        """
+        clamped_x = torch.clamp(x, min=eps)
+        return torch.log(clamped_x)
+
+    @staticmethod
+    def inverse_sigmod(x):
+        r"""
+        Inverse of the sigmoid function
+        """
+        if not torch.is_tensor(x):
+            x = torch.tensor(x)
+        return OverflowUtils.log_clamped(x / (1.0 - x))
+
+    @staticmethod
+    def inverse_softplus(x):
+        r"""
+        Inverse of the softplus function
+        """
+        if not torch.is_tensor(x):
+            x = torch.tensor(x)
+        return OverflowUtils.log_clamped(torch.exp(x) - 1.0)
+
+    @staticmethod
+    def logsumexp(x, dim):
+        r"""
+        Differentiable LogSumExp: Does not creates nan gradients
+            when all the inputs are -inf yeilds 0 gradients.
+        Args:
+            x : torch.Tensor -  The input tensor
+            dim: int - The dimension on which the log sum exp has to be applied
+        """
+
+        m, _ = x.max(dim=dim)
+        mask = m == -float("inf")
+        s = (x - m.masked_fill_(mask, 0).unsqueeze(dim=dim)).exp().sum(dim=dim)
+        return s.masked_fill_(mask, 1).log() + m.masked_fill_(mask, -float("inf"))
+
+    @staticmethod
+    def double_pad(list_of_different_shape_tensors):
+        r"""
+        Pads the list of tensors in 2 dimensions
+        """
+        second_dim_lens = [len(a) for a in [i[0] for i in list_of_different_shape_tensors]]
+        second_dim_max = max(second_dim_lens)
+        padded_x = [F.pad(x, (0, second_dim_max - len(x[0]))) for x in list_of_different_shape_tensors]
+        return nn.utils.rnn.pad_sequence(padded_x, batch_first=True)

TTS/tts/layers/overflow/decoder.py

@@ -0,0 +1,81 @@
+import torch
+from torch import nn
+
+from TTS.tts.layers.glow_tts.decoder import Decoder as GlowDecoder
+from TTS.tts.utils.helpers import sequence_mask
+
+
+class Decoder(nn.Module):
+    """Uses glow decoder with some modifications.
+    ::
+
+        Squeeze -> ActNorm -> InvertibleConv1x1 -> AffineCoupling -> Unsqueeze
+
+    Args:
+        in_channels (int): channels of input tensor.
+        hidden_channels (int): hidden decoder channels.
+        kernel_size (int): Coupling block kernel size. (Wavenet filter kernel size.)
+        dilation_rate (int): rate to increase dilation by each layer in a decoder block.
+        num_flow_blocks (int): number of decoder blocks.
+        num_coupling_layers (int): number coupling layers. (number of wavenet layers.)
+        dropout_p (float): wavenet dropout rate.
+        sigmoid_scale (bool): enable/disable sigmoid scaling in coupling layer.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        num_flow_blocks,
+        num_coupling_layers,
+        dropout_p=0.0,
+        num_splits=4,
+        num_squeeze=2,
+        sigmoid_scale=False,
+        c_in_channels=0,
+    ):
+        super().__init__()
+
+        self.glow_decoder = GlowDecoder(
+            in_channels,
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            num_flow_blocks,
+            num_coupling_layers,
+            dropout_p,
+            num_splits,
+            num_squeeze,
+            sigmoid_scale,
+            c_in_channels,
+        )
+        self.n_sqz = num_squeeze
+
+    def forward(self, x, x_len, g=None, reverse=False):
+        """
+        Input shapes:
+            - x:  :math:`[B, C, T]`
+            - x_len :math:`[B]`
+            - g: :math:`[B, C]`
+
+        Output shapes:
+            - x:  :math:`[B, C, T]`
+            - x_len :math:`[B]`
+            - logget_tot :math:`[B]`
+        """
+        x, x_len, x_max_len = self.preprocess(x, x_len, x_len.max())
+        x_mask = torch.unsqueeze(sequence_mask(x_len, x_max_len), 1).to(x.dtype)
+        x, logdet_tot = self.glow_decoder(x, x_mask, g, reverse)
+        return x, x_len, logdet_tot
+
+    def preprocess(self, y, y_lengths, y_max_length):
+        if y_max_length is not None:
+            y_max_length = torch.div(y_max_length, self.n_sqz, rounding_mode="floor") * self.n_sqz
+            y = y[:, :, :y_max_length]
+        y_lengths = torch.div(y_lengths, self.n_sqz, rounding_mode="floor") * self.n_sqz
+        return y, y_lengths, y_max_length
+
+    def store_inverse(self):
+        self.glow_decoder.store_inverse()

TTS/tts/layers/overflow/neural_hmm.py

@@ -0,0 +1,553 @@
+from typing import List
+
+import torch
+import torch.distributions as tdist
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+from TTS.tts.layers.overflow.common_layers import Outputnet, OverflowUtils
+from TTS.tts.layers.tacotron.common_layers import Prenet
+from TTS.tts.utils.helpers import sequence_mask
+
+
+class NeuralHMM(nn.Module):
+    """Autoregressive left to right HMM model primarily used in "Neural HMMs are all you need (for high-quality attention-free TTS)"
+
+    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.
+
+    Args:
+        frame_channels (int): Output dimension to generate.
+        ar_order (int): Autoregressive order of the model. In ablations of Neural HMM it was found that more autoregression while giving more variation hurts naturalness of the synthesised audio.
+        deterministic_transition (bool): deterministic duration generation based on duration quantiles as defiend in "S. Ronanki, O. Watts, S. King, and G. E. Henter, “Medianbased generation of synthetic speech durations using a nonparametric approach,” in Proc. SLT, 2016.". Defaults to True.
+        encoder_dim (int): Channels of encoder input and character embedding tensors. Defaults to 512.
+        prenet_type (str): `original` or `bn`. `original` sets the default Prenet and `bn` uses Batch Normalization version of the Prenet.
+        prenet_dim (int): Dimension of the Prenet.
+        prenet_n_layers (int): Number of layers in the Prenet.
+        prenet_dropout (float): Dropout probability of the Prenet.
+        prenet_dropout_at_inference (bool): If True, dropout is applied at inference time.
+        memory_rnn_dim (int): Size of the memory RNN to process output of prenet.
+        outputnet_size (List[int]): Size of the output network inside the neural HMM.
+        flat_start_params (dict): Parameters for the flat start initialization of the neural HMM.
+        std_floor (float): Floor value for the standard deviation of the neural HMM. Prevents model cheating by putting point mass and getting infinite likelihood at any datapoint.
+        use_grad_checkpointing (bool, optional): Use gradient checkpointing to save memory. Defaults to True.
+    """
+
+    def __init__(
+        self,
+        frame_channels: int,
+        ar_order: int,
+        deterministic_transition: bool,
+        encoder_dim: int,
+        prenet_type: str,
+        prenet_dim: int,
+        prenet_n_layers: int,
+        prenet_dropout: float,
+        prenet_dropout_at_inference: bool,
+        memory_rnn_dim: int,
+        outputnet_size: List[int],
+        flat_start_params: dict,
+        std_floor: float,
+        use_grad_checkpointing: bool = True,
+    ):
+        super().__init__()
+
+        self.frame_channels = frame_channels
+        self.ar_order = ar_order
+        self.deterministic_transition = deterministic_transition
+        self.prenet_dim = prenet_dim
+        self.memory_rnn_dim = memory_rnn_dim
+        self.use_grad_checkpointing = use_grad_checkpointing
+
+        self.transition_model = TransitionModel()
+        self.emission_model = EmissionModel()
+
+        assert ar_order > 0, f"AR order must be greater than 0 provided {ar_order}"
+
+        self.ar_order = ar_order
+        self.prenet = Prenet(
+            in_features=frame_channels * ar_order,
+            prenet_type=prenet_type,
+            prenet_dropout=prenet_dropout,
+            dropout_at_inference=prenet_dropout_at_inference,
+            out_features=[self.prenet_dim for _ in range(prenet_n_layers)],
+            bias=False,
+        )
+        self.memory_rnn = nn.LSTMCell(input_size=prenet_dim, hidden_size=memory_rnn_dim)
+        self.output_net = Outputnet(
+            encoder_dim, memory_rnn_dim, frame_channels, outputnet_size, flat_start_params, std_floor
+        )
+        self.register_buffer("go_tokens", torch.zeros(ar_order, 1))
+
+    def forward(self, inputs, inputs_len, mels, mel_lens):
+        r"""HMM forward algorithm for training uses logarithmic version of Rabiner (1989) forward algorithm.
+
+        Args:
+            inputs (torch.FloatTensor): Encoder outputs
+            inputs_len (torch.LongTensor): Encoder output lengths
+            mels (torch.FloatTensor): Mel inputs
+            mel_lens (torch.LongTensor): Length of mel inputs
+
+        Shapes:
+            - inputs: (B, T, D_out_enc)
+            - inputs_len: (B)
+            - mels: (B, D_mel, T_mel)
+            - mel_lens: (B)
+
+        Returns:
+            log_prob (torch.FloatTensor): Log probability of the sequence
+        """
+        # Get dimensions of inputs
+        batch_size, N, _ = inputs.shape
+        T_max = torch.max(mel_lens)
+        mels = mels.permute(0, 2, 1)
+
+        # Intialize forward algorithm
+        log_state_priors = self._initialize_log_state_priors(inputs)
+        log_c, log_alpha_scaled, transition_matrix, means = self._initialize_forward_algorithm_variables(mels, N)
+
+        # Initialize autoregression elements
+        ar_inputs = self._add_go_token(mels)
+        h_memory, c_memory = self._init_lstm_states(batch_size, self.memory_rnn_dim, mels)
+
+        for t in range(T_max):
+            # Process Autoregression
+            h_memory, c_memory = self._process_ar_timestep(t, ar_inputs, h_memory, c_memory)
+            # Get mean, std and transition vector from decoder for this timestep
+            # Note: Gradient checkpointing currently doesn't works with multiple gpus inside a loop
+            if self.use_grad_checkpointing and self.training:
+                mean, std, transition_vector = checkpoint(self.output_net, h_memory, inputs)
+            else:
+                mean, std, transition_vector = self.output_net(h_memory, inputs)
+
+            if t == 0:
+                log_alpha_temp = log_state_priors + self.emission_model(mels[:, 0], mean, std, inputs_len)
+            else:
+                log_alpha_temp = self.emission_model(mels[:, t], mean, std, inputs_len) + self.transition_model(
+                    log_alpha_scaled[:, t - 1, :], transition_vector, inputs_len
+                )
+            log_c[:, t] = torch.logsumexp(log_alpha_temp, dim=1)
+            log_alpha_scaled[:, t, :] = log_alpha_temp - log_c[:, t].unsqueeze(1)
+            transition_matrix[:, t] = transition_vector  # needed for absorption state calculation
+
+            # Save for plotting
+            means.append(mean.detach())
+
+        log_c, log_alpha_scaled = self._mask_lengths(mel_lens, log_c, log_alpha_scaled)
+
+        sum_final_log_c = self.get_absorption_state_scaling_factor(
+            mel_lens, log_alpha_scaled, inputs_len, transition_matrix
+        )
+
+        log_probs = torch.sum(log_c, dim=1) + sum_final_log_c
+
+        return log_probs, log_alpha_scaled, transition_matrix, means
+
+    @staticmethod
+    def _mask_lengths(mel_lens, log_c, log_alpha_scaled):
+        """
+        Mask the lengths of the forward variables so that the variable lenghts
+        do not contribute in the loss calculation
+        Args:
+            mel_inputs (torch.FloatTensor): (batch, T, frame_channels)
+            mel_inputs_lengths (torch.IntTensor): (batch)
+            log_c (torch.FloatTensor): (batch, T)
+        Returns:
+            log_c (torch.FloatTensor) : scaled probabilities (batch, T)
+            log_alpha_scaled (torch.FloatTensor): forward probabilities (batch, T, N)
+        """
+        mask_log_c = sequence_mask(mel_lens)
+        log_c = log_c * mask_log_c
+        mask_log_alpha_scaled = mask_log_c.unsqueeze(2)
+        log_alpha_scaled = log_alpha_scaled * mask_log_alpha_scaled
+        return log_c, log_alpha_scaled
+
+    def _process_ar_timestep(
+        self,
+        t,
+        ar_inputs,
+        h_memory,
+        c_memory,
+    ):
+        """
+        Process autoregression in timestep
+        1. At a specific t timestep
+        2. Perform data dropout if applied (we did not use it)
+        3. Run the autoregressive frame through the prenet (has dropout)
+        4. Run the prenet output through the post prenet rnn
+
+        Args:
+            t (int): mel-spec timestep
+            ar_inputs (torch.FloatTensor): go-token appended mel-spectrograms
+                - shape: (b, D_out, T_out)
+            h_post_prenet (torch.FloatTensor): previous timestep rnn hidden state
+                - shape: (b, memory_rnn_dim)
+            c_post_prenet (torch.FloatTensor): previous timestep rnn cell state
+                - shape: (b, memory_rnn_dim)
+
+        Returns:
+            h_post_prenet (torch.FloatTensor): rnn hidden state of the current timestep
+            c_post_prenet (torch.FloatTensor): rnn cell state of the current timestep
+        """
+        prenet_input = ar_inputs[:, t : t + self.ar_order].flatten(1)
+        memory_inputs = self.prenet(prenet_input)
+        h_memory, c_memory = self.memory_rnn(memory_inputs, (h_memory, c_memory))
+        return h_memory, c_memory
+
+    def _add_go_token(self, mel_inputs):
+        """Append the go token to create the autoregressive input
+        Args:
+            mel_inputs (torch.FloatTensor): (batch_size, T, n_mel_channel)
+        Returns:
+            ar_inputs (torch.FloatTensor): (batch_size, T, n_mel_channel)
+        """
+        batch_size, T, _ = mel_inputs.shape
+        go_tokens = self.go_tokens.unsqueeze(0).expand(batch_size, self.ar_order, self.frame_channels)
+        ar_inputs = torch.cat((go_tokens, mel_inputs), dim=1)[:, :T]
+        return ar_inputs
+
+    @staticmethod
+    def _initialize_forward_algorithm_variables(mel_inputs, N):
+        r"""Initialize placeholders for forward algorithm variables, to use a stable
+                version we will use log_alpha_scaled and the scaling constant
+
+        Args:
+            mel_inputs (torch.FloatTensor): (b, T_max, frame_channels)
+            N (int): number of states
+        Returns:
+            log_c (torch.FloatTensor): Scaling constant (b, T_max)
+        """
+        b, T_max, _ = mel_inputs.shape
+        log_alpha_scaled = mel_inputs.new_zeros((b, T_max, N))
+        log_c = mel_inputs.new_zeros(b, T_max)
+        transition_matrix = mel_inputs.new_zeros((b, T_max, N))
+
+        # Saving for plotting later, will not have gradient tapes
+        means = []
+        return log_c, log_alpha_scaled, transition_matrix, means
+
+    @staticmethod
+    def _init_lstm_states(batch_size, hidden_state_dim, device_tensor):
+        r"""
+        Initialize Hidden and Cell states for LSTM Cell
+
+        Args:
+            batch_size (Int): batch size
+            hidden_state_dim (Int): dimensions of the h and c
+            device_tensor (torch.FloatTensor): useful for the device and type
+
+        Returns:
+            (torch.FloatTensor): shape (batch_size, hidden_state_dim)
+                can be hidden state for LSTM
+            (torch.FloatTensor): shape (batch_size, hidden_state_dim)
+                can be the cell state for LSTM
+        """
+        return (
+            device_tensor.new_zeros(batch_size, hidden_state_dim),
+            device_tensor.new_zeros(batch_size, hidden_state_dim),
+        )
+
+    def get_absorption_state_scaling_factor(self, mels_len, log_alpha_scaled, inputs_len, transition_vector):
+        """Returns the final scaling factor of absorption state
+
+        Args:
+            mels_len (torch.IntTensor): Input size of mels to
+                    get the last timestep of log_alpha_scaled
+            log_alpha_scaled (torch.FloatTEnsor): State probabilities
+            text_lengths (torch.IntTensor): length of the states to
+                    mask the values of states lengths
+                (
+                    Useful when the batch has very different lengths,
+                    when the length of an observation is less than
+                    the number of max states, then the log alpha after
+                    the state value is filled with -infs. So we mask
+                    those values so that it only consider the states
+                    which are needed for that length
+                )
+            transition_vector (torch.FloatTensor): transtiion vector for each state per timestep
+
+        Shapes:
+            - mels_len: (batch_size)
+            - log_alpha_scaled: (batch_size, N, T)
+            - text_lengths: (batch_size)
+            - transition_vector: (batch_size, N, T)
+
+        Returns:
+            sum_final_log_c (torch.FloatTensor): (batch_size)
+
+        """
+        N = torch.max(inputs_len)
+        max_inputs_len = log_alpha_scaled.shape[2]
+        state_lengths_mask = sequence_mask(inputs_len, max_len=max_inputs_len)
+
+        last_log_alpha_scaled_index = (
+            (mels_len - 1).unsqueeze(-1).expand(-1, N).unsqueeze(1)
+        )  # Batch X Hidden State Size
+        last_log_alpha_scaled = torch.gather(log_alpha_scaled, 1, last_log_alpha_scaled_index).squeeze(1)
+        last_log_alpha_scaled = last_log_alpha_scaled.masked_fill(~state_lengths_mask, -float("inf"))
+
+        last_transition_vector = torch.gather(transition_vector, 1, last_log_alpha_scaled_index).squeeze(1)
+        last_transition_probability = torch.sigmoid(last_transition_vector)
+        log_probability_of_transitioning = OverflowUtils.log_clamped(last_transition_probability)
+
+        last_transition_probability_index = self.get_mask_for_last_item(inputs_len, inputs_len.device)
+        log_probability_of_transitioning = log_probability_of_transitioning.masked_fill(
+            ~last_transition_probability_index, -float("inf")
+        )
+        final_log_c = last_log_alpha_scaled + log_probability_of_transitioning
+
+        # If the length of the mel is less than the number of states it will select the -inf values leading to nan gradients
+        # Ideally, we should clean the dataset otherwise this is a little hack uncomment the line below
+        final_log_c = final_log_c.clamp(min=torch.finfo(final_log_c.dtype).min)
+
+        sum_final_log_c = torch.logsumexp(final_log_c, dim=1)
+        return sum_final_log_c
+
+    @staticmethod
+    def get_mask_for_last_item(lengths, device, out_tensor=None):
+        """Returns n-1 mask for the last item in the sequence.
+
+        Args:
+            lengths (torch.IntTensor): lengths in a batch
+            device (str, optional): Defaults to "cpu".
+            out_tensor (torch.Tensor, optional): uses the memory of a specific tensor.
+                Defaults to None.
+
+        Returns:
+            - Shape: :math:`(b, max_len)`
+        """
+        max_len = torch.max(lengths).item()
+        ids = (
+            torch.arange(0, max_len, device=device) if out_tensor is None else torch.arange(0, max_len, out=out_tensor)
+        )
+        mask = ids == lengths.unsqueeze(1) - 1
+        return mask
+
+    @torch.inference_mode()
+    def inference(
+        self,
+        inputs: torch.FloatTensor,
+        input_lens: torch.LongTensor,
+        sampling_temp: float,
+        max_sampling_time: int,
+        duration_threshold: float,
+    ):
+        """Inference from autoregressive neural HMM
+
+        Args:
+            inputs (torch.FloatTensor): input states
+                - shape: :math:`(b, T, d)`
+            input_lens (torch.LongTensor): input state lengths
+                - shape: :math:`(b)`
+            sampling_temp (float): sampling temperature
+            max_sampling_temp (int): max sampling temperature
+            duration_threshold (float): duration threshold to switch to next state
+                - Use this to change the spearking rate of the synthesised audio
+        """
+
+        b = inputs.shape[0]
+        outputs = {
+            "hmm_outputs": [],
+            "hmm_outputs_len": [],
+            "alignments": [],
+            "input_parameters": [],
+            "output_parameters": [],
+        }
+        for i in range(b):
+            neural_hmm_outputs, states_travelled, input_parameters, output_parameters = self.sample(
+                inputs[i : i + 1], input_lens[i], sampling_temp, max_sampling_time, duration_threshold
+            )
+
+            outputs["hmm_outputs"].append(neural_hmm_outputs)
+            outputs["hmm_outputs_len"].append(neural_hmm_outputs.shape[0])
+            outputs["alignments"].append(states_travelled)
+            outputs["input_parameters"].append(input_parameters)
+            outputs["output_parameters"].append(output_parameters)
+
+        outputs["hmm_outputs"] = nn.utils.rnn.pad_sequence(outputs["hmm_outputs"], batch_first=True)
+        outputs["hmm_outputs_len"] = torch.tensor(
+            outputs["hmm_outputs_len"], dtype=input_lens.dtype, device=input_lens.device
+        )
+        return outputs
+
+    @torch.inference_mode()
+    def sample(self, inputs, input_lens, sampling_temp, max_sampling_time, duration_threshold):
+        """Samples an output from the parameter models
+
+        Args:
+            inputs (torch.FloatTensor): input states
+                - shape: :math:`(1, T, d)`
+            input_lens (torch.LongTensor): input state lengths
+                - shape: :math:`(1)`
+            sampling_temp (float): sampling temperature
+            max_sampling_time (int): max sampling time
+            duration_threshold (float): duration threshold to switch to next state
+
+        Returns:
+            outputs (torch.FloatTensor): Output Observations
+                - Shape: :math:`(T, output_dim)`
+            states_travelled (list[int]): Hidden states travelled
+                - Shape: :math:`(T)`
+            input_parameters (list[torch.FloatTensor]): Input parameters
+            output_parameters (list[torch.FloatTensor]): Output parameters
+        """
+        states_travelled, outputs, t = [], [], 0
+
+        # Sample initial state
+        current_state = 0
+        states_travelled.append(current_state)
+
+        # Prepare autoregression
+        prenet_input = self.go_tokens.unsqueeze(0).expand(1, self.ar_order, self.frame_channels)
+        h_memory, c_memory = self._init_lstm_states(1, self.memory_rnn_dim, prenet_input)
+
+        input_parameter_values = []
+        output_parameter_values = []
+        quantile = 1
+        while True:
+            memory_input = self.prenet(prenet_input.flatten(1).unsqueeze(0))
+            # will be 1 while sampling
+            h_memory, c_memory = self.memory_rnn(memory_input.squeeze(0), (h_memory, c_memory))
+
+            z_t = inputs[:, current_state].unsqueeze(0)  # Add fake time dimension
+            mean, std, transition_vector = self.output_net(h_memory, z_t)
+
+            transition_probability = torch.sigmoid(transition_vector.flatten())
+            staying_probability = torch.sigmoid(-transition_vector.flatten())
+
+            # Save for plotting
+            input_parameter_values.append([prenet_input, current_state])
+            output_parameter_values.append([mean, std, transition_probability])
+
+            x_t = self.emission_model.sample(mean, std, sampling_temp=sampling_temp)
+
+            # Prepare autoregressive input for next iteration
+            prenet_input = torch.cat((prenet_input, x_t), dim=1)[:, 1:]
+
+            outputs.append(x_t.flatten())
+
+            transition_matrix = torch.cat((staying_probability, transition_probability))
+            quantile *= staying_probability
+            if not self.deterministic_transition:
+                switch = transition_matrix.multinomial(1)[0].item()
+            else:
+                switch = quantile < duration_threshold
+
+            if switch:
+                current_state += 1
+                quantile = 1
+
+            states_travelled.append(current_state)
+
+            if (current_state == input_lens) or (max_sampling_time and t == max_sampling_time - 1):
+                break
+
+            t += 1
+
+        return (
+            torch.stack(outputs, dim=0),
+            F.one_hot(input_lens.new_tensor(states_travelled)),
+            input_parameter_values,
+            output_parameter_values,
+        )
+
+    @staticmethod
+    def _initialize_log_state_priors(text_embeddings):
+        """Creates the log pi in forward algorithm.
+
+        Args:
+            text_embeddings (torch.FloatTensor): used to create the log pi
+                    on current device
+
+        Shapes:
+            - text_embeddings: (B, T, D_out_enc)
+        """
+        N = text_embeddings.shape[1]
+        log_state_priors = text_embeddings.new_full([N], -float("inf"))
+        log_state_priors[0] = 0.0
+        return log_state_priors
+
+
+class TransitionModel(nn.Module):
+    """Transition Model of the HMM, it represents the probability of transitioning
+    form current state to all other states"""
+
+    def forward(self, log_alpha_scaled, transition_vector, inputs_len):  # pylint: disable=no-self-use
+        r"""
+        product of the past state with transitional probabilities in log space
+
+        Args:
+            log_alpha_scaled (torch.Tensor): Multiply previous timestep's alphas by
+                        transition matrix (in log domain)
+                - shape: (batch size, N)
+            transition_vector (torch.tensor): transition vector for each state
+                - shape: (N)
+            inputs_len (int tensor): Lengths of states in a batch
+                - shape: (batch)
+
+        Returns:
+            out (torch.FloatTensor): log probability of transitioning to each state
+        """
+        transition_p = torch.sigmoid(transition_vector)
+        staying_p = torch.sigmoid(-transition_vector)
+
+        log_staying_probability = OverflowUtils.log_clamped(staying_p)
+        log_transition_probability = OverflowUtils.log_clamped(transition_p)
+
+        staying = log_alpha_scaled + log_staying_probability
+        leaving = log_alpha_scaled + log_transition_probability
+        leaving = leaving.roll(1, dims=1)
+        leaving[:, 0] = -float("inf")
+        inputs_len_mask = sequence_mask(inputs_len)
+        out = OverflowUtils.logsumexp(torch.stack((staying, leaving), dim=2), dim=2)
+        out = out.masked_fill(~inputs_len_mask, -float("inf"))  # There are no states to contribute to the loss
+        return out
+
+
+class EmissionModel(nn.Module):
+    """Emission Model of the HMM, it represents the probability of
+    emitting an observation based on the current state"""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.distribution_function: tdist.Distribution = tdist.normal.Normal
+
+    def sample(self, means, stds, sampling_temp):
+        return self.distribution_function(means, stds * sampling_temp).sample() if sampling_temp > 0 else means
+
+    def forward(self, x_t, means, stds, state_lengths):
+        r"""Calculates the log probability of the the given data (x_t)
+            being observed from states with given means and stds
+        Args:
+            x_t (float tensor) : observation at current time step
+                - shape: (batch, feature_dim)
+            means (float tensor): means of the distributions of hidden states
+                - shape: (batch, hidden_state, feature_dim)
+            stds (float tensor): standard deviations of the distributions of the hidden states
+                - shape: (batch, hidden_state, feature_dim)
+            state_lengths (int tensor): Lengths of states in a batch
+                - shape: (batch)
+
+        Returns:
+            out (float tensor): observation log likelihoods,
+                                    expressing the probability of an observation
+                being generated from a state i
+                shape: (batch, hidden_state)
+        """
+        emission_dists = self.distribution_function(means, stds)
+        out = emission_dists.log_prob(x_t.unsqueeze(1))
+        state_lengths_mask = sequence_mask(state_lengths).unsqueeze(2)
+        out = torch.sum(out * state_lengths_mask, dim=2)
+        return out

TTS/tts/layers/overflow/plotting_utils.py

@@ -0,0 +1,79 @@
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+
+def validate_numpy_array(value: Any):
+    r"""
+    Validates the input and makes sure it returns a numpy array (i.e on CPU)
+
+    Args:
+        value (Any): the input value
+
+    Raises:
+        TypeError: if the value is not a numpy array or torch tensor
+
+    Returns:
+        np.ndarray: numpy array of the value
+    """
+    if isinstance(value, np.ndarray):
+        pass
+    elif isinstance(value, list):
+        value = np.array(value)
+    elif torch.is_tensor(value):
+        value = value.cpu().numpy()
+    else:
+        raise TypeError("Value must be a numpy array, a torch tensor or a list")
+
+    return value
+
+
+def get_spec_from_most_probable_state(log_alpha_scaled, means, decoder=None):
+    """Get the most probable state means from the log_alpha_scaled.
+
+    Args:
+        log_alpha_scaled (torch.Tensor): Log alpha scaled values.
+            - Shape: :math:`(T, N)`
+        means (torch.Tensor): Means of the states.
+            - Shape: :math:`(N, T, D_out)`
+        decoder (torch.nn.Module): Decoder module to decode the latent to melspectrogram. Defaults to None.
+    """
+    max_state_numbers = torch.max(log_alpha_scaled, dim=1)[1]
+    max_len = means.shape[0]
+    n_mel_channels = means.shape[2]
+    max_state_numbers = max_state_numbers.unsqueeze(1).unsqueeze(1).expand(max_len, 1, n_mel_channels)
+    means = torch.gather(means, 1, max_state_numbers).squeeze(1).to(log_alpha_scaled.dtype)
+    if decoder is not None:
+        mel = (
+            decoder(means.T.unsqueeze(0), torch.tensor([means.shape[0]], device=means.device), reverse=True)[0]
+            .squeeze(0)
+            .T
+        )
+    else:
+        mel = means
+    return mel
+
+
+def plot_transition_probabilities_to_numpy(states, transition_probabilities, output_fig=False):
+    """Generates trainsition probabilities plot for the states and the probability of transition.
+
+    Args:
+        states (torch.IntTensor): the states
+        transition_probabilities (torch.FloatTensor): the transition probabilities
+    """
+    states = validate_numpy_array(states)
+    transition_probabilities = validate_numpy_array(transition_probabilities)
+
+    fig, ax = plt.subplots(figsize=(30, 3))
+    ax.plot(transition_probabilities, "o")
+    ax.set_title("Transition probability of state")
+    ax.set_xlabel("hidden state")
+    ax.set_ylabel("probability")
+    ax.set_xticks([i for i in range(len(transition_probabilities))])  # pylint: disable=unnecessary-comprehension
+    ax.set_xticklabels([int(x) for x in states], rotation=90)
+    plt.tight_layout()
+    if not output_fig:
+        plt.close()
+    return fig

TTS/tts/layers/tacotron/attentions.py

@@ -0,0 +1,486 @@
+import torch
+from scipy.stats import betabinom
+from torch import nn
+from torch.nn import functional as F
+
+from TTS.tts.layers.tacotron.common_layers import Linear
+
+
+class LocationLayer(nn.Module):
+    """Layers for Location Sensitive Attention
+
+    Args:
+        attention_dim (int): number of channels in the input tensor.
+        attention_n_filters (int, optional): number of filters in convolution. Defaults to 32.
+        attention_kernel_size (int, optional): kernel size of convolution filter. Defaults to 31.
+    """
+
+    def __init__(self, attention_dim, attention_n_filters=32, attention_kernel_size=31):
+        super().__init__()
+        self.location_conv1d = nn.Conv1d(
+            in_channels=2,
+            out_channels=attention_n_filters,
+            kernel_size=attention_kernel_size,
+            stride=1,
+            padding=(attention_kernel_size - 1) // 2,
+            bias=False,
+        )
+        self.location_dense = Linear(attention_n_filters, attention_dim, bias=False, init_gain="tanh")
+
+    def forward(self, attention_cat):
+        """
+        Shapes:
+            attention_cat: [B, 2, C]
+        """
+        processed_attention = self.location_conv1d(attention_cat)
+        processed_attention = self.location_dense(processed_attention.transpose(1, 2))
+        return processed_attention
+
+
+class GravesAttention(nn.Module):
+    """Graves Attention as is ref1 with updates from ref2.
+    ref1: https://arxiv.org/abs/1910.10288
+    ref2: https://arxiv.org/pdf/1906.01083.pdf
+
+    Args:
+        query_dim (int): number of channels in query tensor.
+        K (int): number of Gaussian heads to be used for computing attention.
+    """
+
+    COEF = 0.3989422917366028  # numpy.sqrt(1/(2*numpy.pi))
+
+    def __init__(self, query_dim, K):
+        super().__init__()
+        self._mask_value = 1e-8
+        self.K = K
+        # self.attention_alignment = 0.05
+        self.eps = 1e-5
+        self.J = None
+        self.N_a = nn.Sequential(
+            nn.Linear(query_dim, query_dim, bias=True), nn.ReLU(), nn.Linear(query_dim, 3 * K, bias=True)
+        )
+        self.attention_weights = None
+        self.mu_prev = None
+        self.init_layers()
+
+    def init_layers(self):
+        torch.nn.init.constant_(self.N_a[2].bias[(2 * self.K) : (3 * self.K)], 1.0)  # bias mean
+        torch.nn.init.constant_(self.N_a[2].bias[self.K : (2 * self.K)], 10)  # bias std
+
+    def init_states(self, inputs):
+        if self.J is None or inputs.shape[1] + 1 > self.J.shape[-1]:
+            self.J = torch.arange(0, inputs.shape[1] + 2.0).to(inputs.device) + 0.5
+        self.attention_weights = torch.zeros(inputs.shape[0], inputs.shape[1]).to(inputs.device)
+        self.mu_prev = torch.zeros(inputs.shape[0], self.K).to(inputs.device)
+
+    # pylint: disable=R0201
+    # pylint: disable=unused-argument
+    def preprocess_inputs(self, inputs):
+        return None
+
+    def forward(self, query, inputs, processed_inputs, mask):
+        """
+        Shapes:
+            query: [B, C_attention_rnn]
+            inputs: [B, T_in, C_encoder]
+            processed_inputs: place_holder
+            mask: [B, T_in]
+        """
+        gbk_t = self.N_a(query)
+        gbk_t = gbk_t.view(gbk_t.size(0), -1, self.K)
+
+        # attention model parameters
+        # each B x K
+        g_t = gbk_t[:, 0, :]
+        b_t = gbk_t[:, 1, :]
+        k_t = gbk_t[:, 2, :]
+
+        # dropout to decorrelate attention heads
+        g_t = torch.nn.functional.dropout(g_t, p=0.5, training=self.training)
+
+        # attention GMM parameters
+        sig_t = torch.nn.functional.softplus(b_t) + self.eps
+
+        mu_t = self.mu_prev + torch.nn.functional.softplus(k_t)
+        g_t = torch.softmax(g_t, dim=-1) + self.eps
+
+        j = self.J[: inputs.size(1) + 1]
+
+        # attention weights
+        phi_t = g_t.unsqueeze(-1) * (1 / (1 + torch.sigmoid((mu_t.unsqueeze(-1) - j) / sig_t.unsqueeze(-1))))
+
+        # discritize attention weights
+        alpha_t = torch.sum(phi_t, 1)
+        alpha_t = alpha_t[:, 1:] - alpha_t[:, :-1]
+        alpha_t[alpha_t == 0] = 1e-8
+
+        # apply masking
+        if mask is not None:
+            alpha_t.data.masked_fill_(~mask, self._mask_value)
+
+        context = torch.bmm(alpha_t.unsqueeze(1), inputs).squeeze(1)
+        self.attention_weights = alpha_t
+        self.mu_prev = mu_t
+        return context
+
+
+class OriginalAttention(nn.Module):
+    """Bahdanau Attention with various optional modifications.
+    - Location sensitive attnetion: https://arxiv.org/abs/1712.05884
+    - Forward Attention: https://arxiv.org/abs/1807.06736 + state masking at inference
+    - Using sigmoid instead of softmax normalization
+    - Attention windowing at inference time
+
+    Note:
+        Location Sensitive Attention extends the additive attention mechanism
+    to use cumulative attention weights from previous decoder time steps with the current time step features.
+
+        Forward attention computes most probable monotonic alignment. The modified attention probabilities at each
+    timestep are computed recursively by the forward algorithm.
+
+        Transition agent in the forward attention explicitly gates the attention mechanism whether to move forward or
+    stay at each decoder timestep.
+
+        Attention windowing is a inductive prior that prevents the model from attending to previous and future timesteps
+    beyond a certain window.
+
+    Args:
+        query_dim (int): number of channels in the query tensor.
+        embedding_dim (int): number of channels in the vakue tensor. In general, the value tensor is the output of the encoder layer.
+        attention_dim (int): number of channels of the inner attention layers.
+        location_attention (bool): enable/disable location sensitive attention.
+        attention_location_n_filters (int): number of location attention filters.
+        attention_location_kernel_size (int): filter size of location attention convolution layer.
+        windowing (int): window size for attention windowing. if it is 5, for computing the attention, it only considers the time steps [(t-5), ..., (t+5)] of the input.
+        norm (str): normalization method applied to the attention weights. 'softmax' or 'sigmoid'
+        forward_attn (bool): enable/disable forward attention.
+        trans_agent (bool): enable/disable transition agent in the forward attention.
+        forward_attn_mask (int): enable/disable an explicit masking in forward attention. It is useful to set at especially inference time.
+    """
+
+    # Pylint gets confused by PyTorch conventions here
+    # pylint: disable=attribute-defined-outside-init
+    def __init__(
+        self,
+        query_dim,
+        embedding_dim,
+        attention_dim,
+        location_attention,
+        attention_location_n_filters,
+        attention_location_kernel_size,
+        windowing,
+        norm,
+        forward_attn,
+        trans_agent,
+        forward_attn_mask,
+    ):
+        super().__init__()
+        self.query_layer = Linear(query_dim, attention_dim, bias=False, init_gain="tanh")
+        self.inputs_layer = Linear(embedding_dim, attention_dim, bias=False, init_gain="tanh")
+        self.v = Linear(attention_dim, 1, bias=True)
+        if trans_agent:
+            self.ta = nn.Linear(query_dim + embedding_dim, 1, bias=True)
+        if location_attention:
+            self.location_layer = LocationLayer(
+                attention_dim,
+                attention_location_n_filters,
+                attention_location_kernel_size,
+            )
+        self._mask_value = -float("inf")
+        self.windowing = windowing
+        self.win_idx = None
+        self.norm = norm
+        self.forward_attn = forward_attn
+        self.trans_agent = trans_agent
+        self.forward_attn_mask = forward_attn_mask
+        self.location_attention = location_attention
+
+    def init_win_idx(self):
+        self.win_idx = -1
+        self.win_back = 2
+        self.win_front = 6
+
+    def init_forward_attn(self, inputs):
+        B = inputs.shape[0]
+        T = inputs.shape[1]
+        self.alpha = torch.cat([torch.ones([B, 1]), torch.zeros([B, T])[:, :-1] + 1e-7], dim=1).to(inputs.device)
+        self.u = (0.5 * torch.ones([B, 1])).to(inputs.device)
+
+    def init_location_attention(self, inputs):
+        B = inputs.size(0)
+        T = inputs.size(1)
+        self.attention_weights_cum = torch.zeros([B, T], device=inputs.device)
+
+    def init_states(self, inputs):
+        B = inputs.size(0)
+        T = inputs.size(1)
+        self.attention_weights = torch.zeros([B, T], device=inputs.device)
+        if self.location_attention:
+            self.init_location_attention(inputs)
+        if self.forward_attn:
+            self.init_forward_attn(inputs)
+        if self.windowing:
+            self.init_win_idx()
+
+    def preprocess_inputs(self, inputs):
+        return self.inputs_layer(inputs)
+
+    def update_location_attention(self, alignments):
+        self.attention_weights_cum += alignments
+
+    def get_location_attention(self, query, processed_inputs):
+        attention_cat = torch.cat((self.attention_weights.unsqueeze(1), self.attention_weights_cum.unsqueeze(1)), dim=1)
+        processed_query = self.query_layer(query.unsqueeze(1))
+        processed_attention_weights = self.location_layer(attention_cat)
+        energies = self.v(torch.tanh(processed_query + processed_attention_weights + processed_inputs))
+        energies = energies.squeeze(-1)
+        return energies, processed_query
+
+    def get_attention(self, query, processed_inputs):
+        processed_query = self.query_layer(query.unsqueeze(1))
+        energies = self.v(torch.tanh(processed_query + processed_inputs))
+        energies = energies.squeeze(-1)
+        return energies, processed_query
+
+    def apply_windowing(self, attention, inputs):
+        back_win = self.win_idx - self.win_back
+        front_win = self.win_idx + self.win_front
+        if back_win > 0:
+            attention[:, :back_win] = -float("inf")
+        if front_win < inputs.shape[1]:
+            attention[:, front_win:] = -float("inf")
+        # this is a trick to solve a special problem.
+        # but it does not hurt.
+        if self.win_idx == -1:
+            attention[:, 0] = attention.max()
+        # Update the window
+        self.win_idx = torch.argmax(attention, 1).long()[0].item()
+        return attention
+
+    def apply_forward_attention(self, alignment):
+        # forward attention
+        fwd_shifted_alpha = F.pad(self.alpha[:, :-1].clone().to(alignment.device), (1, 0, 0, 0))
+        # compute transition potentials
+        alpha = ((1 - self.u) * self.alpha + self.u * fwd_shifted_alpha + 1e-8) * alignment
+        # force incremental alignment
+        if not self.training and self.forward_attn_mask:
+            _, n = fwd_shifted_alpha.max(1)
+            val, _ = alpha.max(1)
+            for b in range(alignment.shape[0]):
+                alpha[b, n[b] + 3 :] = 0
+                alpha[b, : (n[b] - 1)] = 0  # ignore all previous states to prevent repetition.
+                alpha[b, (n[b] - 2)] = 0.01 * val[b]  # smoothing factor for the prev step
+        # renormalize attention weights
+        alpha = alpha / alpha.sum(dim=1, keepdim=True)
+        return alpha
+
+    def forward(self, query, inputs, processed_inputs, mask):
+        """
+        shapes:
+            query: [B, C_attn_rnn]
+            inputs: [B, T_en, D_en]
+            processed_inputs: [B, T_en, D_attn]
+            mask: [B, T_en]
+        """
+        if self.location_attention:
+            attention, _ = self.get_location_attention(query, processed_inputs)
+        else:
+            attention, _ = self.get_attention(query, processed_inputs)
+        # apply masking
+        if mask is not None:
+            attention.data.masked_fill_(~mask, self._mask_value)
+        # apply windowing - only in eval mode
+        if not self.training and self.windowing:
+            attention = self.apply_windowing(attention, inputs)
+
+        # normalize attention values
+        if self.norm == "softmax":
+            alignment = torch.softmax(attention, dim=-1)
+        elif self.norm == "sigmoid":
+            alignment = torch.sigmoid(attention) / torch.sigmoid(attention).sum(dim=1, keepdim=True)
+        else:
+            raise ValueError("Unknown value for attention norm type")
+
+        if self.location_attention:
+            self.update_location_attention(alignment)
+
+        # apply forward attention if enabled
+        if self.forward_attn:
+            alignment = self.apply_forward_attention(alignment)
+            self.alpha = alignment
+
+        context = torch.bmm(alignment.unsqueeze(1), inputs)
+        context = context.squeeze(1)
+        self.attention_weights = alignment
+
+        # compute transition agent
+        if self.forward_attn and self.trans_agent:
+            ta_input = torch.cat([context, query.squeeze(1)], dim=-1)
+            self.u = torch.sigmoid(self.ta(ta_input))
+        return context
+
+
+class MonotonicDynamicConvolutionAttention(nn.Module):
+    """Dynamic convolution attention from
+    https://arxiv.org/pdf/1910.10288.pdf
+
+
+    query -> linear -> tanh -> linear ->|
+                                        |                                            mask values
+                                        v                                              |    |
+               atten_w(t-1) -|-> conv1d_dynamic -> linear -|-> tanh -> + -> softmax -> * -> * -> context
+                             |-> conv1d_static  -> linear -|           |
+                             |-> conv1d_prior   -> log ----------------|
+
+    query: attention rnn output.
+
+    Note:
+        Dynamic convolution attention is an alternation of the location senstive attention with
+    dynamically computed convolution filters from the previous attention scores and a set of
+    constraints to keep the attention alignment diagonal.
+        DCA is sensitive to mixed precision training and might cause instable training.
+
+    Args:
+        query_dim (int): number of channels in the query tensor.
+        embedding_dim (int): number of channels in the value tensor.
+        static_filter_dim (int): number of channels in the convolution layer computing the static filters.
+        static_kernel_size (int): kernel size for the convolution layer computing the static filters.
+        dynamic_filter_dim (int): number of channels in the convolution layer computing the dynamic filters.
+        dynamic_kernel_size (int): kernel size for the convolution layer computing the dynamic filters.
+        prior_filter_len (int, optional): [description]. Defaults to 11 from the paper.
+        alpha (float, optional): [description]. Defaults to 0.1 from the paper.
+        beta (float, optional): [description]. Defaults to 0.9 from the paper.
+    """
+
+    def __init__(
+        self,
+        query_dim,
+        embedding_dim,  # pylint: disable=unused-argument
+        attention_dim,
+        static_filter_dim,
+        static_kernel_size,
+        dynamic_filter_dim,
+        dynamic_kernel_size,
+        prior_filter_len=11,
+        alpha=0.1,
+        beta=0.9,
+    ):
+        super().__init__()
+        self._mask_value = 1e-8
+        self.dynamic_filter_dim = dynamic_filter_dim
+        self.dynamic_kernel_size = dynamic_kernel_size
+        self.prior_filter_len = prior_filter_len
+        self.attention_weights = None
+        # setup key and query layers
+        self.query_layer = nn.Linear(query_dim, attention_dim)
+        self.key_layer = nn.Linear(attention_dim, dynamic_filter_dim * dynamic_kernel_size, bias=False)
+        self.static_filter_conv = nn.Conv1d(
+            1,
+            static_filter_dim,
+            static_kernel_size,
+            padding=(static_kernel_size - 1) // 2,
+            bias=False,
+        )
+        self.static_filter_layer = nn.Linear(static_filter_dim, attention_dim, bias=False)
+        self.dynamic_filter_layer = nn.Linear(dynamic_filter_dim, attention_dim)
+        self.v = nn.Linear(attention_dim, 1, bias=False)
+
+        prior = betabinom.pmf(range(prior_filter_len), prior_filter_len - 1, alpha, beta)
+        self.register_buffer("prior", torch.FloatTensor(prior).flip(0))
+
+    # pylint: disable=unused-argument
+    def forward(self, query, inputs, processed_inputs, mask):
+        """
+        query: [B, C_attn_rnn]
+        inputs: [B, T_en, D_en]
+        processed_inputs: place holder.
+        mask: [B, T_en]
+        """
+        # compute prior filters
+        prior_filter = F.conv1d(
+            F.pad(self.attention_weights.unsqueeze(1), (self.prior_filter_len - 1, 0)), self.prior.view(1, 1, -1)
+        )
+        prior_filter = torch.log(prior_filter.clamp_min_(1e-6)).squeeze(1)
+        G = self.key_layer(torch.tanh(self.query_layer(query)))
+        # compute dynamic filters
+        dynamic_filter = F.conv1d(
+            self.attention_weights.unsqueeze(0),
+            G.view(-1, 1, self.dynamic_kernel_size),
+            padding=(self.dynamic_kernel_size - 1) // 2,
+            groups=query.size(0),
+        )
+        dynamic_filter = dynamic_filter.view(query.size(0), self.dynamic_filter_dim, -1).transpose(1, 2)
+        # compute static filters
+        static_filter = self.static_filter_conv(self.attention_weights.unsqueeze(1)).transpose(1, 2)
+        alignment = (
+            self.v(
+                torch.tanh(self.static_filter_layer(static_filter) + self.dynamic_filter_layer(dynamic_filter))
+            ).squeeze(-1)
+            + prior_filter
+        )
+        # compute attention weights
+        attention_weights = F.softmax(alignment, dim=-1)
+        # apply masking
+        if mask is not None:
+            attention_weights.data.masked_fill_(~mask, self._mask_value)
+        self.attention_weights = attention_weights
+        # compute context
+        context = torch.bmm(attention_weights.unsqueeze(1), inputs).squeeze(1)
+        return context
+
+    def preprocess_inputs(self, inputs):  # pylint: disable=no-self-use
+        return None
+
+    def init_states(self, inputs):
+        B = inputs.size(0)
+        T = inputs.size(1)
+        self.attention_weights = torch.zeros([B, T], device=inputs.device)
+        self.attention_weights[:, 0] = 1.0
+
+
+def init_attn(
+    attn_type,
+    query_dim,
+    embedding_dim,
+    attention_dim,
+    location_attention,
+    attention_location_n_filters,
+    attention_location_kernel_size,
+    windowing,
+    norm,
+    forward_attn,
+    trans_agent,
+    forward_attn_mask,
+    attn_K,
+):
+    if attn_type == "original":
+        return OriginalAttention(
+            query_dim,
+            embedding_dim,
+            attention_dim,
+            location_attention,
+            attention_location_n_filters,
+            attention_location_kernel_size,
+            windowing,
+            norm,
+            forward_attn,
+            trans_agent,
+            forward_attn_mask,
+        )
+    if attn_type == "graves":
+        return GravesAttention(query_dim, attn_K)
+    if attn_type == "dynamic_convolution":
+        return MonotonicDynamicConvolutionAttention(
+            query_dim,
+            embedding_dim,
+            attention_dim,
+            static_filter_dim=8,
+            static_kernel_size=21,
+            dynamic_filter_dim=8,
+            dynamic_kernel_size=21,
+            prior_filter_len=11,
+            alpha=0.1,
+            beta=0.9,
+        )
+
+    raise RuntimeError(f" [!] Given Attention Type '{attn_type}' is not exist.")

TTS/tts/layers/tacotron/capacitron_layers.py

@@ -0,0 +1,205 @@
+import torch
+from torch import nn
+from torch.distributions.multivariate_normal import MultivariateNormal as MVN
+from torch.nn import functional as F
+
+
+class CapacitronVAE(nn.Module):
+    """Effective Use of Variational Embedding Capacity for prosody transfer.
+
+    See https://arxiv.org/abs/1906.03402"""
+
+    def __init__(
+        self,
+        num_mel,
+        capacitron_VAE_embedding_dim,
+        encoder_output_dim=256,
+        reference_encoder_out_dim=128,
+        speaker_embedding_dim=None,
+        text_summary_embedding_dim=None,
+    ):
+        super().__init__()
+        # Init distributions
+        self.prior_distribution = MVN(
+            torch.zeros(capacitron_VAE_embedding_dim), torch.eye(capacitron_VAE_embedding_dim)
+        )
+        self.approximate_posterior_distribution = None
+        # define output ReferenceEncoder dim to the capacitron_VAE_embedding_dim
+        self.encoder = ReferenceEncoder(num_mel, out_dim=reference_encoder_out_dim)
+
+        # Init beta, the lagrange-like term for the KL distribution
+        self.beta = torch.nn.Parameter(torch.log(torch.exp(torch.Tensor([1.0])) - 1), requires_grad=True)
+        mlp_input_dimension = reference_encoder_out_dim
+
+        if text_summary_embedding_dim is not None:
+            self.text_summary_net = TextSummary(text_summary_embedding_dim, encoder_output_dim=encoder_output_dim)
+            mlp_input_dimension += text_summary_embedding_dim
+        if speaker_embedding_dim is not None:
+            # TODO: Test a multispeaker model!
+            mlp_input_dimension += speaker_embedding_dim
+        self.post_encoder_mlp = PostEncoderMLP(mlp_input_dimension, capacitron_VAE_embedding_dim)
+
+    def forward(self, reference_mel_info=None, text_info=None, speaker_embedding=None):
+        # Use reference
+        if reference_mel_info is not None:
+            reference_mels = reference_mel_info[0]  # [batch_size, num_frames, num_mels]
+            mel_lengths = reference_mel_info[1]  # [batch_size]
+            enc_out = self.encoder(reference_mels, mel_lengths)
+
+            # concat speaker_embedding and/or text summary embedding
+            if text_info is not None:
+                text_inputs = text_info[0]  # [batch_size, num_characters, num_embedding]
+                input_lengths = text_info[1]
+                text_summary_out = self.text_summary_net(text_inputs, input_lengths).to(reference_mels.device)
+                enc_out = torch.cat([enc_out, text_summary_out], dim=-1)
+            if speaker_embedding is not None:
+                speaker_embedding = torch.squeeze(speaker_embedding)
+                enc_out = torch.cat([enc_out, speaker_embedding], dim=-1)
+
+            # Feed the output of the ref encoder and information about text/speaker into
+            # an MLP to produce the parameteres for the approximate poterior distributions
+            mu, sigma = self.post_encoder_mlp(enc_out)
+            # convert to cpu because prior_distribution was created on cpu
+            mu = mu.cpu()
+            sigma = sigma.cpu()
+
+            # Sample from the posterior: z ~ q(z|x)
+            self.approximate_posterior_distribution = MVN(mu, torch.diag_embed(sigma))
+            VAE_embedding = self.approximate_posterior_distribution.rsample()
+        # Infer from the model, bypasses encoding
+        else:
+            # Sample from the prior: z ~ p(z)
+            VAE_embedding = self.prior_distribution.sample().unsqueeze(0)
+
+        # reshape to [batch_size, 1, capacitron_VAE_embedding_dim]
+        return VAE_embedding.unsqueeze(1), self.approximate_posterior_distribution, self.prior_distribution, self.beta
+
+
+class ReferenceEncoder(nn.Module):
+    """NN module creating a fixed size prosody embedding from a spectrogram.
+
+    inputs: mel spectrograms [batch_size, num_spec_frames, num_mel]
+    outputs: [batch_size, embedding_dim]
+    """
+
+    def __init__(self, num_mel, out_dim):
+        super().__init__()
+        self.num_mel = num_mel
+        filters = [1] + [32, 32, 64, 64, 128, 128]
+        num_layers = len(filters) - 1
+        convs = [
+            nn.Conv2d(
+                in_channels=filters[i], out_channels=filters[i + 1], kernel_size=(3, 3), stride=(2, 2), padding=(2, 2)
+            )
+            for i in range(num_layers)
+        ]
+        self.convs = nn.ModuleList(convs)
+        self.training = False
+        self.bns = nn.ModuleList([nn.BatchNorm2d(num_features=filter_size) for filter_size in filters[1:]])
+
+        post_conv_height = self.calculate_post_conv_height(num_mel, 3, 2, 2, num_layers)
+        self.recurrence = nn.LSTM(
+            input_size=filters[-1] * post_conv_height, hidden_size=out_dim, batch_first=True, bidirectional=False
+        )
+
+    def forward(self, inputs, input_lengths):
+        batch_size = inputs.size(0)
+        x = inputs.view(batch_size, 1, -1, self.num_mel)  # [batch_size, num_channels==1, num_frames, num_mel]
+        valid_lengths = input_lengths.float()  # [batch_size]
+        for conv, bn in zip(self.convs, self.bns):
+            x = conv(x)
+            x = bn(x)
+            x = F.relu(x)
+
+            # Create the post conv width mask based on the valid lengths of the output of the convolution.
+            # The valid lengths for the output of a convolution on varying length inputs is
+            # ceil(input_length/stride) + 1 for stride=3 and padding=2
+            # For example (kernel_size=3, stride=2, padding=2):
+            # 0 0 x x x x x 0 0 -> Input = 5, 0 is zero padding, x is valid values coming from padding=2 in conv2d
+            # _____
+            #   x _____
+            #       x _____
+            #           x  ____
+            #               x
+            # x x x x -> Output valid length = 4
+            # Since every example in te batch is zero padded and therefore have separate valid_lengths,
+            # we need to mask off all the values AFTER the valid length for each example in the batch.
+            # Otherwise, the convolutions create noise and a lot of not real information
+            valid_lengths = (valid_lengths / 2).float()
+            valid_lengths = torch.ceil(valid_lengths).to(dtype=torch.int64) + 1  # 2 is stride -- size: [batch_size]
+            post_conv_max_width = x.size(2)
+
+            mask = torch.arange(post_conv_max_width).to(inputs.device).expand(
+                len(valid_lengths), post_conv_max_width
+            ) < valid_lengths.unsqueeze(1)
+            mask = mask.expand(1, 1, -1, -1).transpose(2, 0).transpose(-1, 2)  # [batch_size, 1, post_conv_max_width, 1]
+            x = x * mask
+
+        x = x.transpose(1, 2)
+        # x: 4D tensor [batch_size, post_conv_width,
+        #               num_channels==128, post_conv_height]
+
+        post_conv_width = x.size(1)
+        x = x.contiguous().view(batch_size, post_conv_width, -1)
+        # x: 3D tensor [batch_size, post_conv_width,
+        #               num_channels*post_conv_height]
+
+        # Routine for fetching the last valid output of a dynamic LSTM with varying input lengths and padding
+        post_conv_input_lengths = valid_lengths
+        packed_seqs = nn.utils.rnn.pack_padded_sequence(
+            x, post_conv_input_lengths.tolist(), batch_first=True, enforce_sorted=False
+        )  # dynamic rnn sequence padding
+        self.recurrence.flatten_parameters()
+        _, (ht, _) = self.recurrence(packed_seqs)
+        last_output = ht[-1]
+
+        return last_output.to(inputs.device)  # [B, 128]
+
+    @staticmethod
+    def calculate_post_conv_height(height, kernel_size, stride, pad, n_convs):
+        """Height of spec after n convolutions with fixed kernel/stride/pad."""
+        for _ in range(n_convs):
+            height = (height - kernel_size + 2 * pad) // stride + 1
+        return height
+
+
+class TextSummary(nn.Module):
+    def __init__(self, embedding_dim, encoder_output_dim):
+        super().__init__()
+        self.lstm = nn.LSTM(
+            encoder_output_dim,  # text embedding dimension from the text encoder
+            embedding_dim,  # fixed length output summary the lstm creates from the input
+            batch_first=True,
+            bidirectional=False,
+        )
+
+    def forward(self, inputs, input_lengths):
+        # Routine for fetching the last valid output of a dynamic LSTM with varying input lengths and padding
+        packed_seqs = nn.utils.rnn.pack_padded_sequence(
+            inputs, input_lengths.tolist(), batch_first=True, enforce_sorted=False
+        )  # dynamic rnn sequence padding
+        self.lstm.flatten_parameters()
+        _, (ht, _) = self.lstm(packed_seqs)
+        last_output = ht[-1]
+        return last_output
+
+
+class PostEncoderMLP(nn.Module):
+    def __init__(self, input_size, hidden_size):
+        super().__init__()
+        self.hidden_size = hidden_size
+        modules = [
+            nn.Linear(input_size, hidden_size),  # Hidden Layer
+            nn.Tanh(),
+            nn.Linear(hidden_size, hidden_size * 2),
+        ]  # Output layer twice the size for mean and variance
+        self.net = nn.Sequential(*modules)
+        self.softplus = nn.Softplus()
+
+    def forward(self, _input):
+        mlp_output = self.net(_input)
+        # The mean parameter is unconstrained
+        mu = mlp_output[:, : self.hidden_size]
+        # The standard deviation must be positive. Parameterise with a softplus
+        sigma = self.softplus(mlp_output[:, self.hidden_size :])
+        return mu, sigma

TTS/tts/layers/tacotron/common_layers.py

@@ -0,0 +1,119 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class Linear(nn.Module):
+    """Linear layer with a specific initialization.
+
+    Args:
+        in_features (int): number of channels in the input tensor.
+        out_features (int): number of channels in the output tensor.
+        bias (bool, optional): enable/disable bias in the layer. Defaults to True.
+        init_gain (str, optional): method to compute the gain in the weight initializtion based on the nonlinear activation used afterwards. Defaults to 'linear'.
+    """
+
+    def __init__(self, in_features, out_features, bias=True, init_gain="linear"):
+        super().__init__()
+        self.linear_layer = torch.nn.Linear(in_features, out_features, bias=bias)
+        self._init_w(init_gain)
+
+    def _init_w(self, init_gain):
+        torch.nn.init.xavier_uniform_(self.linear_layer.weight, gain=torch.nn.init.calculate_gain(init_gain))
+
+    def forward(self, x):
+        return self.linear_layer(x)
+
+
+class LinearBN(nn.Module):
+    """Linear layer with Batch Normalization.
+
+    x -> linear -> BN -> o
+
+    Args:
+        in_features (int): number of channels in the input tensor.
+        out_features (int ): number of channels in the output tensor.
+        bias (bool, optional): enable/disable bias in the linear layer. Defaults to True.
+        init_gain (str, optional): method to set the gain for weight initialization. Defaults to 'linear'.
+    """
+
+    def __init__(self, in_features, out_features, bias=True, init_gain="linear"):
+        super().__init__()
+        self.linear_layer = torch.nn.Linear(in_features, out_features, bias=bias)
+        self.batch_normalization = nn.BatchNorm1d(out_features, momentum=0.1, eps=1e-5)
+        self._init_w(init_gain)
+
+    def _init_w(self, init_gain):
+        torch.nn.init.xavier_uniform_(self.linear_layer.weight, gain=torch.nn.init.calculate_gain(init_gain))
+
+    def forward(self, x):
+        """
+        Shapes:
+            x: [T, B, C] or [B, C]
+        """
+        out = self.linear_layer(x)
+        if len(out.shape) == 3:
+            out = out.permute(1, 2, 0)
+        out = self.batch_normalization(out)
+        if len(out.shape) == 3:
+            out = out.permute(2, 0, 1)
+        return out
+
+
+class Prenet(nn.Module):
+    """Tacotron specific Prenet with an optional Batch Normalization.
+
+    Note:
+        Prenet with BN improves the model performance significantly especially
+    if it is enabled after learning a diagonal attention alignment with the original
+    prenet. However, if the target dataset is high quality then it also works from
+    the start. It is also suggested to disable dropout if BN is in use.
+
+        prenet_type == "original"
+            x -> [linear -> ReLU -> Dropout]xN -> o
+
+        prenet_type == "bn"
+            x -> [linear -> BN -> ReLU -> Dropout]xN -> o
+
+    Args:
+        in_features (int): number of channels in the input tensor and the inner layers.
+        prenet_type (str, optional): prenet type "original" or "bn". Defaults to "original".
+        prenet_dropout (bool, optional): dropout rate. Defaults to True.
+        dropout_at_inference (bool, optional): use dropout at inference. It leads to a better quality for some models.
+        out_features (list, optional): List of output channels for each prenet block.
+            It also defines number of the prenet blocks based on the length of argument list.
+            Defaults to [256, 256].
+        bias (bool, optional): enable/disable bias in prenet linear layers. Defaults to True.
+    """
+
+    # pylint: disable=dangerous-default-value
+    def __init__(
+        self,
+        in_features,
+        prenet_type="original",
+        prenet_dropout=True,
+        dropout_at_inference=False,
+        out_features=[256, 256],
+        bias=True,
+    ):
+        super().__init__()
+        self.prenet_type = prenet_type
+        self.prenet_dropout = prenet_dropout
+        self.dropout_at_inference = dropout_at_inference
+        in_features = [in_features] + out_features[:-1]
+        if prenet_type == "bn":
+            self.linear_layers = nn.ModuleList(
+                [LinearBN(in_size, out_size, bias=bias) for (in_size, out_size) in zip(in_features, out_features)]
+            )
+        elif prenet_type == "original":
+            self.linear_layers = nn.ModuleList(
+                [Linear(in_size, out_size, bias=bias) for (in_size, out_size) in zip(in_features, out_features)]
+            )
+
+    def forward(self, x):
+        for linear in self.linear_layers:
+            if self.prenet_dropout:
+                x = F.dropout(F.relu(linear(x)), p=0.5, training=self.training or self.dropout_at_inference)
+            else:
+                x = F.relu(linear(x))
+        return x

TTS/tts/layers/tacotron/gst_layers.py

@@ -0,0 +1,149 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class GST(nn.Module):
+    """Global Style Token Module for factorizing prosody in speech.
+
+    See https://arxiv.org/pdf/1803.09017"""
+
+    def __init__(self, num_mel, num_heads, num_style_tokens, gst_embedding_dim, embedded_speaker_dim=None):
+        super().__init__()
+        self.encoder = ReferenceEncoder(num_mel, gst_embedding_dim)
+        self.style_token_layer = StyleTokenLayer(num_heads, num_style_tokens, gst_embedding_dim, embedded_speaker_dim)
+
+    def forward(self, inputs, speaker_embedding=None):
+        enc_out = self.encoder(inputs)
+        # concat speaker_embedding
+        if speaker_embedding is not None:
+            enc_out = torch.cat([enc_out, speaker_embedding], dim=-1)
+        style_embed = self.style_token_layer(enc_out)
+
+        return style_embed
+
+
+class ReferenceEncoder(nn.Module):
+    """NN module creating a fixed size prosody embedding from a spectrogram.
+
+    inputs: mel spectrograms [batch_size, num_spec_frames, num_mel]
+    outputs: [batch_size, embedding_dim]
+    """
+
+    def __init__(self, num_mel, embedding_dim):
+        super().__init__()
+        self.num_mel = num_mel
+        filters = [1] + [32, 32, 64, 64, 128, 128]
+        num_layers = len(filters) - 1
+        convs = [
+            nn.Conv2d(
+                in_channels=filters[i], out_channels=filters[i + 1], kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)
+            )
+            for i in range(num_layers)
+        ]
+        self.convs = nn.ModuleList(convs)
+        self.bns = nn.ModuleList([nn.BatchNorm2d(num_features=filter_size) for filter_size in filters[1:]])
+
+        post_conv_height = self.calculate_post_conv_height(num_mel, 3, 2, 1, num_layers)
+        self.recurrence = nn.GRU(
+            input_size=filters[-1] * post_conv_height, hidden_size=embedding_dim // 2, batch_first=True
+        )
+
+    def forward(self, inputs):
+        batch_size = inputs.size(0)
+        x = inputs.view(batch_size, 1, -1, self.num_mel)
+        # x: 4D tensor [batch_size, num_channels==1, num_frames, num_mel]
+        for conv, bn in zip(self.convs, self.bns):
+            x = conv(x)
+            x = bn(x)
+            x = F.relu(x)
+
+        x = x.transpose(1, 2)
+        # x: 4D tensor [batch_size, post_conv_width,
+        #               num_channels==128, post_conv_height]
+        post_conv_width = x.size(1)
+        x = x.contiguous().view(batch_size, post_conv_width, -1)
+        # x: 3D tensor [batch_size, post_conv_width,
+        #               num_channels*post_conv_height]
+        self.recurrence.flatten_parameters()
+        _, out = self.recurrence(x)
+        # out: 3D tensor [seq_len==1, batch_size, encoding_size=128]
+
+        return out.squeeze(0)
+
+    @staticmethod
+    def calculate_post_conv_height(height, kernel_size, stride, pad, n_convs):
+        """Height of spec after n convolutions with fixed kernel/stride/pad."""
+        for _ in range(n_convs):
+            height = (height - kernel_size + 2 * pad) // stride + 1
+        return height
+
+
+class StyleTokenLayer(nn.Module):
+    """NN Module attending to style tokens based on prosody encodings."""
+
+    def __init__(self, num_heads, num_style_tokens, gst_embedding_dim, d_vector_dim=None):
+        super().__init__()
+
+        self.query_dim = gst_embedding_dim // 2
+
+        if d_vector_dim:
+            self.query_dim += d_vector_dim
+
+        self.key_dim = gst_embedding_dim // num_heads
+        self.style_tokens = nn.Parameter(torch.FloatTensor(num_style_tokens, self.key_dim))
+        nn.init.normal_(self.style_tokens, mean=0, std=0.5)
+        self.attention = MultiHeadAttention(
+            query_dim=self.query_dim, key_dim=self.key_dim, num_units=gst_embedding_dim, num_heads=num_heads
+        )
+
+    def forward(self, inputs):
+        batch_size = inputs.size(0)
+        prosody_encoding = inputs.unsqueeze(1)
+        # prosody_encoding: 3D tensor [batch_size, 1, encoding_size==128]
+        tokens = torch.tanh(self.style_tokens).unsqueeze(0).expand(batch_size, -1, -1)
+        # tokens: 3D tensor [batch_size, num tokens, token embedding size]
+        style_embed = self.attention(prosody_encoding, tokens)
+
+        return style_embed
+
+
+class MultiHeadAttention(nn.Module):
+    """
+    input:
+        query --- [N, T_q, query_dim]
+        key --- [N, T_k, key_dim]
+    output:
+        out --- [N, T_q, num_units]
+    """
+
+    def __init__(self, query_dim, key_dim, num_units, num_heads):
+        super().__init__()
+        self.num_units = num_units
+        self.num_heads = num_heads
+        self.key_dim = key_dim
+
+        self.W_query = nn.Linear(in_features=query_dim, out_features=num_units, bias=False)
+        self.W_key = nn.Linear(in_features=key_dim, out_features=num_units, bias=False)
+        self.W_value = nn.Linear(in_features=key_dim, out_features=num_units, bias=False)
+
+    def forward(self, query, key):
+        queries = self.W_query(query)  # [N, T_q, num_units]
+        keys = self.W_key(key)  # [N, T_k, num_units]
+        values = self.W_value(key)
+
+        split_size = self.num_units // self.num_heads
+        queries = torch.stack(torch.split(queries, split_size, dim=2), dim=0)  # [h, N, T_q, num_units/h]
+        keys = torch.stack(torch.split(keys, split_size, dim=2), dim=0)  # [h, N, T_k, num_units/h]
+        values = torch.stack(torch.split(values, split_size, dim=2), dim=0)  # [h, N, T_k, num_units/h]
+
+        # score = softmax(QK^T / (d_k**0.5))
+        scores = torch.matmul(queries, keys.transpose(2, 3))  # [h, N, T_q, T_k]
+        scores = scores / (self.key_dim**0.5)
+        scores = F.softmax(scores, dim=3)
+
+        # out = score * V
+        out = torch.matmul(scores, values)  # [h, N, T_q, num_units/h]
+        out = torch.cat(torch.split(out, 1, dim=0), dim=3).squeeze(0)  # [N, T_q, num_units]
+
+        return out

TTS/tts/layers/tacotron/tacotron.py

@@ -0,0 +1,503 @@
+# coding: utf-8
+# adapted from https://github.com/r9y9/tacotron_pytorch
+
+import torch
+from torch import nn
+
+from .attentions import init_attn
+from .common_layers import Prenet
+
+
+class BatchNormConv1d(nn.Module):
+    r"""A wrapper for Conv1d with BatchNorm. It sets the activation
+    function between Conv and BatchNorm layers. BatchNorm layer
+    is initialized with the TF default values for momentum and eps.
+
+    Args:
+        in_channels: size of each input sample
+        out_channels: size of each output samples
+        kernel_size: kernel size of conv filters
+        stride: stride of conv filters
+        padding: padding of conv filters
+        activation: activation function set b/w Conv1d and BatchNorm
+
+    Shapes:
+        - input: (B, D)
+        - output: (B, D)
+    """
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, activation=None):
+        super().__init__()
+        self.padding = padding
+        self.padder = nn.ConstantPad1d(padding, 0)
+        self.conv1d = nn.Conv1d(
+            in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=0, bias=False
+        )
+        # Following tensorflow's default parameters
+        self.bn = nn.BatchNorm1d(out_channels, momentum=0.99, eps=1e-3)
+        self.activation = activation
+        # self.init_layers()
+
+    def init_layers(self):
+        if isinstance(self.activation, torch.nn.ReLU):
+            w_gain = "relu"
+        elif isinstance(self.activation, torch.nn.Tanh):
+            w_gain = "tanh"
+        elif self.activation is None:
+            w_gain = "linear"
+        else:
+            raise RuntimeError("Unknown activation function")
+        torch.nn.init.xavier_uniform_(self.conv1d.weight, gain=torch.nn.init.calculate_gain(w_gain))
+
+    def forward(self, x):
+        x = self.padder(x)
+        x = self.conv1d(x)
+        x = self.bn(x)
+        if self.activation is not None:
+            x = self.activation(x)
+        return x
+
+
+class Highway(nn.Module):
+    r"""Highway layers as explained in https://arxiv.org/abs/1505.00387
+
+    Args:
+        in_features (int): size of each input sample
+        out_feature (int): size of each output sample
+
+    Shapes:
+        - input: (B, *, H_in)
+        - output: (B, *, H_out)
+    """
+
+    # TODO: Try GLU layer
+    def __init__(self, in_features, out_feature):
+        super().__init__()
+        self.H = nn.Linear(in_features, out_feature)
+        self.H.bias.data.zero_()
+        self.T = nn.Linear(in_features, out_feature)
+        self.T.bias.data.fill_(-1)
+        self.relu = nn.ReLU()
+        self.sigmoid = nn.Sigmoid()
+        # self.init_layers()
+
+    def init_layers(self):
+        torch.nn.init.xavier_uniform_(self.H.weight, gain=torch.nn.init.calculate_gain("relu"))
+        torch.nn.init.xavier_uniform_(self.T.weight, gain=torch.nn.init.calculate_gain("sigmoid"))
+
+    def forward(self, inputs):
+        H = self.relu(self.H(inputs))
+        T = self.sigmoid(self.T(inputs))
+        return H * T + inputs * (1.0 - T)
+
+
+class CBHG(nn.Module):
+    """CBHG module: a recurrent neural network composed of:
+    - 1-d convolution banks
+    - Highway networks + residual connections
+    - Bidirectional gated recurrent units
+
+    Args:
+        in_features (int): sample size
+        K (int): max filter size in conv bank
+        projections (list): conv channel sizes for conv projections
+        num_highways (int): number of highways layers
+
+    Shapes:
+        - input: (B, C, T_in)
+        - output: (B, T_in, C*2)
+    """
+
+    # pylint: disable=dangerous-default-value
+    def __init__(
+        self,
+        in_features,
+        K=16,
+        conv_bank_features=128,
+        conv_projections=[128, 128],
+        highway_features=128,
+        gru_features=128,
+        num_highways=4,
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.conv_bank_features = conv_bank_features
+        self.highway_features = highway_features
+        self.gru_features = gru_features
+        self.conv_projections = conv_projections
+        self.relu = nn.ReLU()
+        # list of conv1d bank with filter size k=1...K
+        # TODO: try dilational layers instead
+        self.conv1d_banks = nn.ModuleList(
+            [
+                BatchNormConv1d(
+                    in_features,
+                    conv_bank_features,
+                    kernel_size=k,
+                    stride=1,
+                    padding=[(k - 1) // 2, k // 2],
+                    activation=self.relu,
+                )
+                for k in range(1, K + 1)
+            ]
+        )
+        # max pooling of conv bank, with padding
+        # TODO: try average pooling OR larger kernel size
+        out_features = [K * conv_bank_features] + conv_projections[:-1]
+        activations = [self.relu] * (len(conv_projections) - 1)
+        activations += [None]
+        # setup conv1d projection layers
+        layer_set = []
+        for in_size, out_size, ac in zip(out_features, conv_projections, activations):
+            layer = BatchNormConv1d(in_size, out_size, kernel_size=3, stride=1, padding=[1, 1], activation=ac)
+            layer_set.append(layer)
+        self.conv1d_projections = nn.ModuleList(layer_set)
+        # setup Highway layers
+        if self.highway_features != conv_projections[-1]:
+            self.pre_highway = nn.Linear(conv_projections[-1], highway_features, bias=False)
+        self.highways = nn.ModuleList([Highway(highway_features, highway_features) for _ in range(num_highways)])
+        # bi-directional GPU layer
+        self.gru = nn.GRU(gru_features, gru_features, 1, batch_first=True, bidirectional=True)
+
+    def forward(self, inputs):
+        # (B, in_features, T_in)
+        x = inputs
+        # (B, hid_features*K, T_in)
+        # Concat conv1d bank outputs
+        outs = []
+        for conv1d in self.conv1d_banks:
+            out = conv1d(x)
+            outs.append(out)
+        x = torch.cat(outs, dim=1)
+        assert x.size(1) == self.conv_bank_features * len(self.conv1d_banks)
+        for conv1d in self.conv1d_projections:
+            x = conv1d(x)
+        x += inputs
+        x = x.transpose(1, 2)
+        if self.highway_features != self.conv_projections[-1]:
+            x = self.pre_highway(x)
+        # Residual connection
+        # TODO: try residual scaling as in Deep Voice 3
+        # TODO: try plain residual layers
+        for highway in self.highways:
+            x = highway(x)
+        # (B, T_in, hid_features*2)
+        # TODO: replace GRU with convolution as in Deep Voice 3
+        self.gru.flatten_parameters()
+        outputs, _ = self.gru(x)
+        return outputs
+
+
+class EncoderCBHG(nn.Module):
+    r"""CBHG module with Encoder specific arguments"""
+
+    def __init__(self):
+        super().__init__()
+        self.cbhg = CBHG(
+            128,
+            K=16,
+            conv_bank_features=128,
+            conv_projections=[128, 128],
+            highway_features=128,
+            gru_features=128,
+            num_highways=4,
+        )
+
+    def forward(self, x):
+        return self.cbhg(x)
+
+
+class Encoder(nn.Module):
+    r"""Stack Prenet and CBHG module for encoder
+    Args:
+        inputs (FloatTensor): embedding features
+
+    Shapes:
+        - inputs: (B, T, D_in)
+        - outputs: (B, T, 128 * 2)
+    """
+
+    def __init__(self, in_features):
+        super().__init__()
+        self.prenet = Prenet(in_features, out_features=[256, 128])
+        self.cbhg = EncoderCBHG()
+
+    def forward(self, inputs):
+        # B x T x prenet_dim
+        outputs = self.prenet(inputs)
+        outputs = self.cbhg(outputs.transpose(1, 2))
+        return outputs
+
+
+class PostCBHG(nn.Module):
+    def __init__(self, mel_dim):
+        super().__init__()
+        self.cbhg = CBHG(
+            mel_dim,
+            K=8,
+            conv_bank_features=128,
+            conv_projections=[256, mel_dim],
+            highway_features=128,
+            gru_features=128,
+            num_highways=4,
+        )
+
+    def forward(self, x):
+        return self.cbhg(x)
+
+
+class Decoder(nn.Module):
+    """Tacotron decoder.
+
+    Args:
+        in_channels (int): number of input channels.
+        frame_channels (int): number of feature frame channels.
+        r (int): number of outputs per time step (reduction rate).
+        memory_size (int): size of the past window. if <= 0 memory_size = r
+        attn_type (string): type of attention used in decoder.
+        attn_windowing (bool): if true, define an attention window centered to maximum
+            attention response. It provides more robust attention alignment especially
+            at interence time.
+        attn_norm (string): attention normalization function. 'sigmoid' or 'softmax'.
+        prenet_type (string): 'original' or 'bn'.
+        prenet_dropout (float): prenet dropout rate.
+        forward_attn (bool): if true, use forward attention method. https://arxiv.org/abs/1807.06736
+        trans_agent (bool): if true, use transition agent. https://arxiv.org/abs/1807.06736
+        forward_attn_mask (bool): if true, mask attention values smaller than a threshold.
+        location_attn (bool): if true, use location sensitive attention.
+        attn_K (int): number of attention heads for GravesAttention.
+        separate_stopnet (bool): if true, detach stopnet input to prevent gradient flow.
+        d_vector_dim (int): size of speaker embedding vector, for multi-speaker training.
+        max_decoder_steps (int): Maximum number of steps allowed for the decoder. Defaults to 500.
+    """
+
+    # Pylint gets confused by PyTorch conventions here
+    # pylint: disable=attribute-defined-outside-init
+
+    def __init__(
+        self,
+        in_channels,
+        frame_channels,
+        r,
+        memory_size,
+        attn_type,
+        attn_windowing,
+        attn_norm,
+        prenet_type,
+        prenet_dropout,
+        forward_attn,
+        trans_agent,
+        forward_attn_mask,
+        location_attn,
+        attn_K,
+        separate_stopnet,
+        max_decoder_steps,
+    ):
+        super().__init__()
+        self.r_init = r
+        self.r = r
+        self.in_channels = in_channels
+        self.max_decoder_steps = max_decoder_steps
+        self.use_memory_queue = memory_size > 0
+        self.memory_size = memory_size if memory_size > 0 else r
+        self.frame_channels = frame_channels
+        self.separate_stopnet = separate_stopnet
+        self.query_dim = 256
+        # memory -> |Prenet| -> processed_memory
+        prenet_dim = frame_channels * self.memory_size if self.use_memory_queue else frame_channels
+        self.prenet = Prenet(prenet_dim, prenet_type, prenet_dropout, out_features=[256, 128])
+        # processed_inputs, processed_memory -> |Attention| -> Attention, attention, RNN_State
+        # attention_rnn generates queries for the attention mechanism
+        self.attention_rnn = nn.GRUCell(in_channels + 128, self.query_dim)
+        self.attention = init_attn(
+            attn_type=attn_type,
+            query_dim=self.query_dim,
+            embedding_dim=in_channels,
+            attention_dim=128,
+            location_attention=location_attn,
+            attention_location_n_filters=32,
+            attention_location_kernel_size=31,
+            windowing=attn_windowing,
+            norm=attn_norm,
+            forward_attn=forward_attn,
+            trans_agent=trans_agent,
+            forward_attn_mask=forward_attn_mask,
+            attn_K=attn_K,
+        )
+        # (processed_memory | attention context) -> |Linear| -> decoder_RNN_input
+        self.project_to_decoder_in = nn.Linear(256 + in_channels, 256)
+        # decoder_RNN_input -> |RNN| -> RNN_state
+        self.decoder_rnns = nn.ModuleList([nn.GRUCell(256, 256) for _ in range(2)])
+        # RNN_state -> |Linear| -> mel_spec
+        self.proj_to_mel = nn.Linear(256, frame_channels * self.r_init)
+        # learn init values instead of zero init.
+        self.stopnet = StopNet(256 + frame_channels * self.r_init)
+
+    def set_r(self, new_r):
+        self.r = new_r
+
+    def _reshape_memory(self, memory):
+        """
+        Reshape the spectrograms for given 'r'
+        """
+        # Grouping multiple frames if necessary
+        if memory.size(-1) == self.frame_channels:
+            memory = memory.view(memory.shape[0], memory.size(1) // self.r, -1)
+        # Time first (T_decoder, B, frame_channels)
+        memory = memory.transpose(0, 1)
+        return memory
+
+    def _init_states(self, inputs):
+        """
+        Initialization of decoder states
+        """
+        B = inputs.size(0)
+        # go frame as zeros matrix
+        if self.use_memory_queue:
+            self.memory_input = torch.zeros(1, device=inputs.device).repeat(B, self.frame_channels * self.memory_size)
+        else:
+            self.memory_input = torch.zeros(1, device=inputs.device).repeat(B, self.frame_channels)
+        # decoder states
+        self.attention_rnn_hidden = torch.zeros(1, device=inputs.device).repeat(B, 256)
+        self.decoder_rnn_hiddens = [
+            torch.zeros(1, device=inputs.device).repeat(B, 256) for idx in range(len(self.decoder_rnns))
+        ]
+        self.context_vec = inputs.data.new(B, self.in_channels).zero_()
+        # cache attention inputs
+        self.processed_inputs = self.attention.preprocess_inputs(inputs)
+
+    def _parse_outputs(self, outputs, attentions, stop_tokens):
+        # Back to batch first
+        attentions = torch.stack(attentions).transpose(0, 1)
+        stop_tokens = torch.stack(stop_tokens).transpose(0, 1)
+        outputs = torch.stack(outputs).transpose(0, 1).contiguous()
+        outputs = outputs.view(outputs.size(0), -1, self.frame_channels)
+        outputs = outputs.transpose(1, 2)
+        return outputs, attentions, stop_tokens
+
+    def decode(self, inputs, mask=None):
+        # Prenet
+        processed_memory = self.prenet(self.memory_input)
+        # Attention RNN
+        self.attention_rnn_hidden = self.attention_rnn(
+            torch.cat((processed_memory, self.context_vec), -1), self.attention_rnn_hidden
+        )
+        self.context_vec = self.attention(self.attention_rnn_hidden, inputs, self.processed_inputs, mask)
+        # Concat RNN output and attention context vector
+        decoder_input = self.project_to_decoder_in(torch.cat((self.attention_rnn_hidden, self.context_vec), -1))
+
+        # Pass through the decoder RNNs
+        for idx, decoder_rnn in enumerate(self.decoder_rnns):
+            self.decoder_rnn_hiddens[idx] = decoder_rnn(decoder_input, self.decoder_rnn_hiddens[idx])
+            # Residual connection
+            decoder_input = self.decoder_rnn_hiddens[idx] + decoder_input
+        decoder_output = decoder_input
+
+        # predict mel vectors from decoder vectors
+        output = self.proj_to_mel(decoder_output)
+        # output = torch.sigmoid(output)
+        # predict stop token
+        stopnet_input = torch.cat([decoder_output, output], -1)
+        if self.separate_stopnet:
+            stop_token = self.stopnet(stopnet_input.detach())
+        else:
+            stop_token = self.stopnet(stopnet_input)
+        output = output[:, : self.r * self.frame_channels]
+        return output, stop_token, self.attention.attention_weights
+
+    def _update_memory_input(self, new_memory):
+        if self.use_memory_queue:
+            if self.memory_size > self.r:
+                # memory queue size is larger than number of frames per decoder iter
+                self.memory_input = torch.cat(
+                    [new_memory, self.memory_input[:, : (self.memory_size - self.r) * self.frame_channels].clone()],
+                    dim=-1,
+                )
+            else:
+                # memory queue size smaller than number of frames per decoder iter
+                self.memory_input = new_memory[:, : self.memory_size * self.frame_channels]
+        else:
+            # use only the last frame prediction
+            # assert new_memory.shape[-1] == self.r * self.frame_channels
+            self.memory_input = new_memory[:, self.frame_channels * (self.r - 1) :]
+
+    def forward(self, inputs, memory, mask):
+        """
+        Args:
+            inputs: Encoder outputs.
+            memory: Decoder memory (autoregression. If None (at eval-time),
+              decoder outputs are used as decoder inputs. If None, it uses the last
+              output as the input.
+            mask: Attention mask for sequence padding.
+
+        Shapes:
+            - inputs: (B, T, D_out_enc)
+            - memory: (B, T_mel, D_mel)
+        """
+        # Run greedy decoding if memory is None
+        memory = self._reshape_memory(memory)
+        outputs = []
+        attentions = []
+        stop_tokens = []
+        t = 0
+        self._init_states(inputs)
+        self.attention.init_states(inputs)
+        while len(outputs) < memory.size(0):
+            if t > 0:
+                new_memory = memory[t - 1]
+                self._update_memory_input(new_memory)
+
+            output, stop_token, attention = self.decode(inputs, mask)
+            outputs += [output]
+            attentions += [attention]
+            stop_tokens += [stop_token.squeeze(1)]
+            t += 1
+        return self._parse_outputs(outputs, attentions, stop_tokens)
+
+    def inference(self, inputs):
+        """
+        Args:
+            inputs: encoder outputs.
+        Shapes:
+            - inputs: batch x time x encoder_out_dim
+        """
+        outputs = []
+        attentions = []
+        stop_tokens = []
+        t = 0
+        self._init_states(inputs)
+        self.attention.init_states(inputs)
+        while True:
+            if t > 0:
+                new_memory = outputs[-1]
+                self._update_memory_input(new_memory)
+            output, stop_token, attention = self.decode(inputs, None)
+            stop_token = torch.sigmoid(stop_token.data)
+            outputs += [output]
+            attentions += [attention]
+            stop_tokens += [stop_token]
+            t += 1
+            if t > inputs.shape[1] / 4 and (stop_token > 0.6 or attention[:, -1].item() > 0.6):
+                break
+            if t > self.max_decoder_steps:
+                print("   | > Decoder stopped with 'max_decoder_steps")
+                break
+        return self._parse_outputs(outputs, attentions, stop_tokens)
+
+
+class StopNet(nn.Module):
+    r"""Stopnet signalling decoder to stop inference.
+    Args:
+        in_features (int): feature dimension of input.
+    """
+
+    def __init__(self, in_features):
+        super().__init__()
+        self.dropout = nn.Dropout(0.1)
+        self.linear = nn.Linear(in_features, 1)
+        torch.nn.init.xavier_uniform_(self.linear.weight, gain=torch.nn.init.calculate_gain("linear"))
+
+    def forward(self, inputs):
+        outputs = self.dropout(inputs)
+        outputs = self.linear(outputs)
+        return outputs

TTS/tts/layers/tacotron/tacotron2.py

@@ -0,0 +1,414 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from .attentions import init_attn
+from .common_layers import Linear, Prenet
+
+
+# pylint: disable=no-value-for-parameter
+# pylint: disable=unexpected-keyword-arg
+class ConvBNBlock(nn.Module):
+    r"""Convolutions with Batch Normalization and non-linear activation.
+
+    Args:
+        in_channels (int): number of input channels.
+        out_channels (int): number of output channels.
+        kernel_size (int): convolution kernel size.
+        activation (str): 'relu', 'tanh', None (linear).
+
+    Shapes:
+        - input: (B, C_in, T)
+        - output: (B, C_out, T)
+    """
+
+    def __init__(self, in_channels, out_channels, kernel_size, activation=None):
+        super().__init__()
+        assert (kernel_size - 1) % 2 == 0
+        padding = (kernel_size - 1) // 2
+        self.convolution1d = nn.Conv1d(in_channels, out_channels, kernel_size, padding=padding)
+        self.batch_normalization = nn.BatchNorm1d(out_channels, momentum=0.1, eps=1e-5)
+        self.dropout = nn.Dropout(p=0.5)
+        if activation == "relu":
+            self.activation = nn.ReLU()
+        elif activation == "tanh":
+            self.activation = nn.Tanh()
+        else:
+            self.activation = nn.Identity()
+
+    def forward(self, x):
+        o = self.convolution1d(x)
+        o = self.batch_normalization(o)
+        o = self.activation(o)
+        o = self.dropout(o)
+        return o
+
+
+class Postnet(nn.Module):
+    r"""Tacotron2 Postnet
+
+    Args:
+        in_out_channels (int): number of output channels.
+
+    Shapes:
+        - input: (B, C_in, T)
+        - output: (B, C_in, T)
+    """
+
+    def __init__(self, in_out_channels, num_convs=5):
+        super().__init__()
+        self.convolutions = nn.ModuleList()
+        self.convolutions.append(ConvBNBlock(in_out_channels, 512, kernel_size=5, activation="tanh"))
+        for _ in range(1, num_convs - 1):
+            self.convolutions.append(ConvBNBlock(512, 512, kernel_size=5, activation="tanh"))
+        self.convolutions.append(ConvBNBlock(512, in_out_channels, kernel_size=5, activation=None))
+
+    def forward(self, x):
+        o = x
+        for layer in self.convolutions:
+            o = layer(o)
+        return o
+
+
+class Encoder(nn.Module):
+    r"""Tacotron2 Encoder
+
+    Args:
+        in_out_channels (int): number of input and output channels.
+
+    Shapes:
+        - input: (B, C_in, T)
+        - output: (B, C_in, T)
+    """
+
+    def __init__(self, in_out_channels=512):
+        super().__init__()
+        self.convolutions = nn.ModuleList()
+        for _ in range(3):
+            self.convolutions.append(ConvBNBlock(in_out_channels, in_out_channels, 5, "relu"))
+        self.lstm = nn.LSTM(
+            in_out_channels, int(in_out_channels / 2), num_layers=1, batch_first=True, bias=True, bidirectional=True
+        )
+        self.rnn_state = None
+
+    def forward(self, x, input_lengths):
+        o = x
+        for layer in self.convolutions:
+            o = layer(o)
+        o = o.transpose(1, 2)
+        o = nn.utils.rnn.pack_padded_sequence(o, input_lengths.cpu(), batch_first=True)
+        self.lstm.flatten_parameters()
+        o, _ = self.lstm(o)
+        o, _ = nn.utils.rnn.pad_packed_sequence(o, batch_first=True)
+        return o
+
+    def inference(self, x):
+        o = x
+        for layer in self.convolutions:
+            o = layer(o)
+        o = o.transpose(1, 2)
+        # self.lstm.flatten_parameters()
+        o, _ = self.lstm(o)
+        return o
+
+
+# adapted from https://github.com/NVIDIA/tacotron2/
+class Decoder(nn.Module):
+    """Tacotron2 decoder. We don't use Zoneout but Dropout between RNN layers.
+
+    Args:
+        in_channels (int): number of input channels.
+        frame_channels (int): number of feature frame channels.
+        r (int): number of outputs per time step (reduction rate).
+        memory_size (int): size of the past window. if <= 0 memory_size = r
+        attn_type (string): type of attention used in decoder.
+        attn_win (bool): if true, define an attention window centered to maximum
+            attention response. It provides more robust attention alignment especially
+            at interence time.
+        attn_norm (string): attention normalization function. 'sigmoid' or 'softmax'.
+        prenet_type (string): 'original' or 'bn'.
+        prenet_dropout (float): prenet dropout rate.
+        forward_attn (bool): if true, use forward attention method. https://arxiv.org/abs/1807.06736
+        trans_agent (bool): if true, use transition agent. https://arxiv.org/abs/1807.06736
+        forward_attn_mask (bool): if true, mask attention values smaller than a threshold.
+        location_attn (bool): if true, use location sensitive attention.
+        attn_K (int): number of attention heads for GravesAttention.
+        separate_stopnet (bool): if true, detach stopnet input to prevent gradient flow.
+        max_decoder_steps (int): Maximum number of steps allowed for the decoder. Defaults to 10000.
+    """
+
+    # Pylint gets confused by PyTorch conventions here
+    # pylint: disable=attribute-defined-outside-init
+    def __init__(
+        self,
+        in_channels,
+        frame_channels,
+        r,
+        attn_type,
+        attn_win,
+        attn_norm,
+        prenet_type,
+        prenet_dropout,
+        forward_attn,
+        trans_agent,
+        forward_attn_mask,
+        location_attn,
+        attn_K,
+        separate_stopnet,
+        max_decoder_steps,
+    ):
+        super().__init__()
+        self.frame_channels = frame_channels
+        self.r_init = r
+        self.r = r
+        self.encoder_embedding_dim = in_channels
+        self.separate_stopnet = separate_stopnet
+        self.max_decoder_steps = max_decoder_steps
+        self.stop_threshold = 0.5
+
+        # model dimensions
+        self.query_dim = 1024
+        self.decoder_rnn_dim = 1024
+        self.prenet_dim = 256
+        self.attn_dim = 128
+        self.p_attention_dropout = 0.1
+        self.p_decoder_dropout = 0.1
+
+        # memory -> |Prenet| -> processed_memory
+        prenet_dim = self.frame_channels
+        self.prenet = Prenet(
+            prenet_dim, prenet_type, prenet_dropout, out_features=[self.prenet_dim, self.prenet_dim], bias=False
+        )
+
+        self.attention_rnn = nn.LSTMCell(self.prenet_dim + in_channels, self.query_dim, bias=True)
+
+        self.attention = init_attn(
+            attn_type=attn_type,
+            query_dim=self.query_dim,
+            embedding_dim=in_channels,
+            attention_dim=128,
+            location_attention=location_attn,
+            attention_location_n_filters=32,
+            attention_location_kernel_size=31,
+            windowing=attn_win,
+            norm=attn_norm,
+            forward_attn=forward_attn,
+            trans_agent=trans_agent,
+            forward_attn_mask=forward_attn_mask,
+            attn_K=attn_K,
+        )
+
+        self.decoder_rnn = nn.LSTMCell(self.query_dim + in_channels, self.decoder_rnn_dim, bias=True)
+
+        self.linear_projection = Linear(self.decoder_rnn_dim + in_channels, self.frame_channels * self.r_init)
+
+        self.stopnet = nn.Sequential(
+            nn.Dropout(0.1),
+            Linear(self.decoder_rnn_dim + self.frame_channels * self.r_init, 1, bias=True, init_gain="sigmoid"),
+        )
+        self.memory_truncated = None
+
+    def set_r(self, new_r):
+        self.r = new_r
+
+    def get_go_frame(self, inputs):
+        B = inputs.size(0)
+        memory = torch.zeros(1, device=inputs.device).repeat(B, self.frame_channels * self.r)
+        return memory
+
+    def _init_states(self, inputs, mask, keep_states=False):
+        B = inputs.size(0)
+        # T = inputs.size(1)
+        if not keep_states:
+            self.query = torch.zeros(1, device=inputs.device).repeat(B, self.query_dim)
+            self.attention_rnn_cell_state = torch.zeros(1, device=inputs.device).repeat(B, self.query_dim)
+            self.decoder_hidden = torch.zeros(1, device=inputs.device).repeat(B, self.decoder_rnn_dim)
+            self.decoder_cell = torch.zeros(1, device=inputs.device).repeat(B, self.decoder_rnn_dim)
+            self.context = torch.zeros(1, device=inputs.device).repeat(B, self.encoder_embedding_dim)
+        self.inputs = inputs
+        self.processed_inputs = self.attention.preprocess_inputs(inputs)
+        self.mask = mask
+
+    def _reshape_memory(self, memory):
+        """
+        Reshape the spectrograms for given 'r'
+        """
+        # Grouping multiple frames if necessary
+        if memory.size(-1) == self.frame_channels:
+            memory = memory.view(memory.shape[0], memory.size(1) // self.r, -1)
+        # Time first (T_decoder, B, frame_channels)
+        memory = memory.transpose(0, 1)
+        return memory
+
+    def _parse_outputs(self, outputs, stop_tokens, alignments):
+        alignments = torch.stack(alignments).transpose(0, 1)
+        stop_tokens = torch.stack(stop_tokens).transpose(0, 1)
+        outputs = torch.stack(outputs).transpose(0, 1).contiguous()
+        outputs = outputs.view(outputs.size(0), -1, self.frame_channels)
+        outputs = outputs.transpose(1, 2)
+        return outputs, stop_tokens, alignments
+
+    def _update_memory(self, memory):
+        if len(memory.shape) == 2:
+            return memory[:, self.frame_channels * (self.r - 1) :]
+        return memory[:, :, self.frame_channels * (self.r - 1) :]
+
+    def decode(self, memory):
+        """
+        shapes:
+           - memory: B x r * self.frame_channels
+        """
+        # self.context: B x D_en
+        # query_input: B x D_en + (r * self.frame_channels)
+        query_input = torch.cat((memory, self.context), -1)
+        # self.query and self.attention_rnn_cell_state : B x D_attn_rnn
+        self.query, self.attention_rnn_cell_state = self.attention_rnn(
+            query_input, (self.query, self.attention_rnn_cell_state)
+        )
+        self.query = F.dropout(self.query, self.p_attention_dropout, self.training)
+        self.attention_rnn_cell_state = F.dropout(
+            self.attention_rnn_cell_state, self.p_attention_dropout, self.training
+        )
+        # B x D_en
+        self.context = self.attention(self.query, self.inputs, self.processed_inputs, self.mask)
+        # B x (D_en + D_attn_rnn)
+        decoder_rnn_input = torch.cat((self.query, self.context), -1)
+        # self.decoder_hidden and self.decoder_cell: B x D_decoder_rnn
+        self.decoder_hidden, self.decoder_cell = self.decoder_rnn(
+            decoder_rnn_input, (self.decoder_hidden, self.decoder_cell)
+        )
+        self.decoder_hidden = F.dropout(self.decoder_hidden, self.p_decoder_dropout, self.training)
+        # B x (D_decoder_rnn + D_en)
+        decoder_hidden_context = torch.cat((self.decoder_hidden, self.context), dim=1)
+        # B x (self.r * self.frame_channels)
+        decoder_output = self.linear_projection(decoder_hidden_context)
+        # B x (D_decoder_rnn + (self.r * self.frame_channels))
+        stopnet_input = torch.cat((self.decoder_hidden, decoder_output), dim=1)
+        if self.separate_stopnet:
+            stop_token = self.stopnet(stopnet_input.detach())
+        else:
+            stop_token = self.stopnet(stopnet_input)
+        # select outputs for the reduction rate self.r
+        decoder_output = decoder_output[:, : self.r * self.frame_channels]
+        return decoder_output, self.attention.attention_weights, stop_token
+
+    def forward(self, inputs, memories, mask):
+        r"""Train Decoder with teacher forcing.
+        Args:
+            inputs: Encoder outputs.
+            memories: Feature frames for teacher-forcing.
+            mask: Attention mask for sequence padding.
+
+        Shapes:
+            - inputs: (B, T, D_out_enc)
+            - memory: (B, T_mel, D_mel)
+            - outputs: (B, T_mel, D_mel)
+            - alignments: (B, T_in, T_out)
+            - stop_tokens: (B, T_out)
+        """
+        memory = self.get_go_frame(inputs).unsqueeze(0)
+        memories = self._reshape_memory(memories)
+        memories = torch.cat((memory, memories), dim=0)
+        memories = self._update_memory(memories)
+        memories = self.prenet(memories)
+
+        self._init_states(inputs, mask=mask)
+        self.attention.init_states(inputs)
+
+        outputs, stop_tokens, alignments = [], [], []
+        while len(outputs) < memories.size(0) - 1:
+            memory = memories[len(outputs)]
+            decoder_output, attention_weights, stop_token = self.decode(memory)
+            outputs += [decoder_output.squeeze(1)]
+            stop_tokens += [stop_token.squeeze(1)]
+            alignments += [attention_weights]
+
+        outputs, stop_tokens, alignments = self._parse_outputs(outputs, stop_tokens, alignments)
+        return outputs, alignments, stop_tokens
+
+    def inference(self, inputs):
+        r"""Decoder inference without teacher forcing and use
+        Stopnet to stop decoder.
+        Args:
+            inputs: Encoder outputs.
+
+        Shapes:
+            - inputs: (B, T, D_out_enc)
+            - outputs: (B, T_mel, D_mel)
+            - alignments: (B, T_in, T_out)
+            - stop_tokens: (B, T_out)
+        """
+        memory = self.get_go_frame(inputs)
+        memory = self._update_memory(memory)
+
+        self._init_states(inputs, mask=None)
+        self.attention.init_states(inputs)
+
+        outputs, stop_tokens, alignments, t = [], [], [], 0
+        while True:
+            memory = self.prenet(memory)
+            decoder_output, alignment, stop_token = self.decode(memory)
+            stop_token = torch.sigmoid(stop_token.data)
+            outputs += [decoder_output.squeeze(1)]
+            stop_tokens += [stop_token]
+            alignments += [alignment]
+
+            if stop_token > self.stop_threshold and t > inputs.shape[0] // 2:
+                break
+            if len(outputs) == self.max_decoder_steps:
+                print(f"   > Decoder stopped with `max_decoder_steps` {self.max_decoder_steps}")
+                break
+
+            memory = self._update_memory(decoder_output)
+            t += 1
+
+        outputs, stop_tokens, alignments = self._parse_outputs(outputs, stop_tokens, alignments)
+
+        return outputs, alignments, stop_tokens
+
+    def inference_truncated(self, inputs):
+        """
+        Preserve decoder states for continuous inference
+        """
+        if self.memory_truncated is None:
+            self.memory_truncated = self.get_go_frame(inputs)
+            self._init_states(inputs, mask=None, keep_states=False)
+        else:
+            self._init_states(inputs, mask=None, keep_states=True)
+
+        self.attention.init_states(inputs)
+        outputs, stop_tokens, alignments, t = [], [], [], 0
+        while True:
+            memory = self.prenet(self.memory_truncated)
+            decoder_output, alignment, stop_token = self.decode(memory)
+            stop_token = torch.sigmoid(stop_token.data)
+            outputs += [decoder_output.squeeze(1)]
+            stop_tokens += [stop_token]
+            alignments += [alignment]
+
+            if stop_token > 0.7:
+                break
+            if len(outputs) == self.max_decoder_steps:
+                print("   | > Decoder stopped with 'max_decoder_steps")
+                break
+
+            self.memory_truncated = decoder_output
+            t += 1
+
+        outputs, stop_tokens, alignments = self._parse_outputs(outputs, stop_tokens, alignments)
+
+        return outputs, alignments, stop_tokens
+
+    def inference_step(self, inputs, t, memory=None):
+        """
+        For debug purposes
+        """
+        if t == 0:
+            memory = self.get_go_frame(inputs)
+            self._init_states(inputs, mask=None)
+
+        memory = self.prenet(memory)
+        decoder_output, stop_token, alignment = self.decode(memory)
+        stop_token = torch.sigmoid(stop_token.data)
+        memory = decoder_output
+        return decoder_output, stop_token, alignment

TTS/tts/layers/tortoise/arch_utils.py

@@ -0,0 +1,433 @@
+import functools
+import math
+import os
+
+import fsspec
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+from transformers import LogitsWarper
+
+from TTS.tts.layers.tortoise.xtransformers import ContinuousTransformerWrapper, RelativePositionBias
+
+
+def zero_module(module):
+    """
+    Zero out the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+class GroupNorm32(nn.GroupNorm):
+    def forward(self, x):
+        return super().forward(x.float()).type(x.dtype)
+
+
+def normalization(channels):
+    """
+    Make a standard normalization layer.
+
+    :param channels: number of input channels.
+    :return: an nn.Module for normalization.
+    """
+    groups = 32
+    if channels <= 16:
+        groups = 8
+    elif channels <= 64:
+        groups = 16
+    while channels % groups != 0:
+        groups = int(groups / 2)
+    assert groups > 2
+    return GroupNorm32(groups, channels)
+
+
+class QKVAttentionLegacy(nn.Module):
+    """
+    A module which performs QKV attention. Matches legacy QKVAttention + input/output heads shaping
+    """
+
+    def __init__(self, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+
+    def forward(self, qkv, mask=None, rel_pos=None):
+        """
+        Apply QKV attention.
+
+        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
+        :return: an [N x (H * C) x T] tensor after attention.
+        """
+        bs, width, length = qkv.shape
+        assert width % (3 * self.n_heads) == 0
+        ch = width // (3 * self.n_heads)
+        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
+        scale = 1 / math.sqrt(math.sqrt(ch))
+        weight = torch.einsum("bct,bcs->bts", q * scale, k * scale)  # More stable with f16 than dividing afterwards
+        if rel_pos is not None:
+            weight = rel_pos(weight.reshape(bs, self.n_heads, weight.shape[-2], weight.shape[-1])).reshape(
+                bs * self.n_heads, weight.shape[-2], weight.shape[-1]
+            )
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        if mask is not None:
+            # The proper way to do this is to mask before the softmax using -inf, but that doesn't work properly on CPUs.
+            mask = mask.repeat(self.n_heads, 1).unsqueeze(1)
+            weight = weight * mask
+        a = torch.einsum("bts,bcs->bct", weight, v)
+
+        return a.reshape(bs, -1, length)
+
+
+class AttentionBlock(nn.Module):
+    """
+    An attention block that allows spatial positions to attend to each other.
+
+    Originally ported from here, but adapted to the N-d case.
+    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
+    """
+
+    def __init__(
+        self,
+        channels,
+        num_heads=1,
+        num_head_channels=-1,
+        do_checkpoint=True,
+        relative_pos_embeddings=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.do_checkpoint = do_checkpoint
+        if num_head_channels == -1:
+            self.num_heads = num_heads
+        else:
+            assert (
+                channels % num_head_channels == 0
+            ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
+            self.num_heads = channels // num_head_channels
+        self.norm = normalization(channels)
+        self.qkv = nn.Conv1d(channels, channels * 3, 1)
+        # split heads before split qkv
+        self.attention = QKVAttentionLegacy(self.num_heads)
+
+        self.proj_out = zero_module(nn.Conv1d(channels, channels, 1))
+        if relative_pos_embeddings:
+            self.relative_pos_embeddings = RelativePositionBias(
+                scale=(channels // self.num_heads) ** 0.5,
+                causal=False,
+                heads=num_heads,
+                num_buckets=32,
+                max_distance=64,
+            )
+        else:
+            self.relative_pos_embeddings = None
+
+    def forward(self, x, mask=None):
+        b, c, *spatial = x.shape
+        x = x.reshape(b, c, -1)
+        qkv = self.qkv(self.norm(x))
+        h = self.attention(qkv, mask, self.relative_pos_embeddings)
+        h = self.proj_out(h)
+        return (x + h).reshape(b, c, *spatial)
+
+
+class Upsample(nn.Module):
+    """
+    An upsampling layer with an optional convolution.
+
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    """
+
+    def __init__(self, channels, use_conv, out_channels=None, factor=4):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.factor = factor
+        if use_conv:
+            ksize = 5
+            pad = 2
+            self.conv = nn.Conv1d(self.channels, self.out_channels, ksize, padding=pad)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        x = F.interpolate(x, scale_factor=self.factor, mode="nearest")
+        if self.use_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    """
+    A downsampling layer with an optional convolution.
+
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    """
+
+    def __init__(self, channels, use_conv, out_channels=None, factor=4, ksize=5, pad=2):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+
+        stride = factor
+        if use_conv:
+            self.op = nn.Conv1d(self.channels, self.out_channels, ksize, stride=stride, padding=pad)
+        else:
+            assert self.channels == self.out_channels
+            self.op = nn.AvgPool1d(kernel_size=stride, stride=stride)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        return self.op(x)
+
+
+class ResBlock(nn.Module):
+    def __init__(
+        self,
+        channels,
+        dropout,
+        out_channels=None,
+        use_conv=False,
+        use_scale_shift_norm=False,
+        up=False,
+        down=False,
+        kernel_size=3,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_scale_shift_norm = use_scale_shift_norm
+        padding = 1 if kernel_size == 3 else 2
+
+        self.in_layers = nn.Sequential(
+            normalization(channels),
+            nn.SiLU(),
+            nn.Conv1d(channels, self.out_channels, kernel_size, padding=padding),
+        )
+
+        self.updown = up or down
+
+        if up:
+            self.h_upd = Upsample(channels, False)
+            self.x_upd = Upsample(channels, False)
+        elif down:
+            self.h_upd = Downsample(channels, False)
+            self.x_upd = Downsample(channels, False)
+        else:
+            self.h_upd = self.x_upd = nn.Identity()
+
+        self.out_layers = nn.Sequential(
+            normalization(self.out_channels),
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            zero_module(nn.Conv1d(self.out_channels, self.out_channels, kernel_size, padding=padding)),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        elif use_conv:
+            self.skip_connection = nn.Conv1d(channels, self.out_channels, kernel_size, padding=padding)
+        else:
+            self.skip_connection = nn.Conv1d(channels, self.out_channels, 1)
+
+    def forward(self, x):
+        if self.updown:
+            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+            h = in_rest(x)
+            h = self.h_upd(h)
+            x = self.x_upd(x)
+            h = in_conv(h)
+        else:
+            h = self.in_layers(x)
+        h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+
+class AudioMiniEncoder(nn.Module):
+    def __init__(
+        self,
+        spec_dim,
+        embedding_dim,
+        base_channels=128,
+        depth=2,
+        resnet_blocks=2,
+        attn_blocks=4,
+        num_attn_heads=4,
+        dropout=0,
+        downsample_factor=2,
+        kernel_size=3,
+    ):
+        super().__init__()
+        self.init = nn.Sequential(nn.Conv1d(spec_dim, base_channels, 3, padding=1))
+        ch = base_channels
+        res = []
+        for l in range(depth):
+            for r in range(resnet_blocks):
+                res.append(ResBlock(ch, dropout, kernel_size=kernel_size))
+            res.append(Downsample(ch, use_conv=True, out_channels=ch * 2, factor=downsample_factor))
+            ch *= 2
+        self.res = nn.Sequential(*res)
+        self.final = nn.Sequential(normalization(ch), nn.SiLU(), nn.Conv1d(ch, embedding_dim, 1))
+        attn = []
+        for a in range(attn_blocks):
+            attn.append(
+                AttentionBlock(
+                    embedding_dim,
+                    num_attn_heads,
+                )
+            )
+        self.attn = nn.Sequential(*attn)
+        self.dim = embedding_dim
+
+    def forward(self, x):
+        h = self.init(x)
+        h = self.res(h)
+        h = self.final(h)
+        h = self.attn(h)
+        return h[:, :, 0]
+
+
+DEFAULT_MEL_NORM_FILE = "https://coqui.gateway.scarf.sh/v0.14.1_models/mel_norms.pth"
+
+
+class TorchMelSpectrogram(nn.Module):
+    def __init__(
+        self,
+        filter_length=1024,
+        hop_length=256,
+        win_length=1024,
+        n_mel_channels=80,
+        mel_fmin=0,
+        mel_fmax=8000,
+        sampling_rate=22050,
+        normalize=False,
+        mel_norm_file=DEFAULT_MEL_NORM_FILE,
+    ):
+        super().__init__()
+        # These are the default tacotron values for the MEL spectrogram.
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.n_mel_channels = n_mel_channels
+        self.mel_fmin = mel_fmin
+        self.mel_fmax = mel_fmax
+        self.sampling_rate = sampling_rate
+        self.mel_stft = torchaudio.transforms.MelSpectrogram(
+            n_fft=self.filter_length,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            power=2,
+            normalized=normalize,
+            sample_rate=self.sampling_rate,
+            f_min=self.mel_fmin,
+            f_max=self.mel_fmax,
+            n_mels=self.n_mel_channels,
+            norm="slaney",
+        )
+        self.mel_norm_file = mel_norm_file
+        if self.mel_norm_file is not None:
+            with fsspec.open(self.mel_norm_file) as f:
+                self.mel_norms = torch.load(f)
+        else:
+            self.mel_norms = None
+
+    def forward(self, inp):
+        if (
+            len(inp.shape) == 3
+        ):  # Automatically squeeze out the channels dimension if it is present (assuming mono-audio)
+            inp = inp.squeeze(1)
+        assert len(inp.shape) == 2
+        self.mel_stft = self.mel_stft.to(inp.device)
+        mel = self.mel_stft(inp)
+        # Perform dynamic range compression
+        mel = torch.log(torch.clamp(mel, min=1e-5))
+        if self.mel_norms is not None:
+            self.mel_norms = self.mel_norms.to(mel.device)
+            mel = mel / self.mel_norms.unsqueeze(0).unsqueeze(-1)
+        return mel
+
+
+class CheckpointedLayer(nn.Module):
+    """
+    Wraps a module. When forward() is called, passes kwargs that require_grad through torch.checkpoint() and bypasses
+    checkpoint for all other args.
+    """
+
+    def __init__(self, wrap):
+        super().__init__()
+        self.wrap = wrap
+
+    def forward(self, x, *args, **kwargs):
+        for k, v in kwargs.items():
+            assert not (isinstance(v, torch.Tensor) and v.requires_grad)  # This would screw up checkpointing.
+        partial = functools.partial(self.wrap, **kwargs)
+        return partial(x, *args)
+
+
+class CheckpointedXTransformerEncoder(nn.Module):
+    """
+    Wraps a ContinuousTransformerWrapper and applies CheckpointedLayer to each layer and permutes from channels-mid
+    to channels-last that XTransformer expects.
+    """
+
+    def __init__(self, needs_permute=True, exit_permute=True, checkpoint=True, **xtransformer_kwargs):
+        super().__init__()
+        self.transformer = ContinuousTransformerWrapper(**xtransformer_kwargs)
+        self.needs_permute = needs_permute
+        self.exit_permute = exit_permute
+
+        if not checkpoint:
+            return
+        for i in range(len(self.transformer.attn_layers.layers)):
+            n, b, r = self.transformer.attn_layers.layers[i]
+            self.transformer.attn_layers.layers[i] = nn.ModuleList([n, CheckpointedLayer(b), r])
+
+    def forward(self, x, **kwargs):
+        if self.needs_permute:
+            x = x.permute(0, 2, 1)
+        h = self.transformer(x, **kwargs)
+        if self.exit_permute:
+            h = h.permute(0, 2, 1)
+        return h
+
+
+class TypicalLogitsWarper(LogitsWarper):
+    def __init__(
+        self,
+        mass: float = 0.9,
+        filter_value: float = -float("Inf"),
+        min_tokens_to_keep: int = 1,
+    ):
+        self.filter_value = filter_value
+        self.mass = mass
+        self.min_tokens_to_keep = min_tokens_to_keep
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        # calculate entropy
+        normalized = torch.nn.functional.log_softmax(scores, dim=-1)
+        p = torch.exp(normalized)
+        ent = -(normalized * p).nansum(-1, keepdim=True)
+
+        # shift and sort
+        shifted_scores = torch.abs((-normalized) - ent)
+        sorted_scores, sorted_indices = torch.sort(shifted_scores, descending=False)
+        sorted_logits = scores.gather(-1, sorted_indices)
+        cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
+
+        # Remove tokens with cumulative mass above the threshold
+        last_ind = (cumulative_probs < self.mass).sum(dim=1)
+        last_ind[last_ind < 0] = 0
+        sorted_indices_to_remove = sorted_scores > sorted_scores.gather(1, last_ind.view(-1, 1))
+        if self.min_tokens_to_keep > 1:
+            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+            sorted_indices_to_remove[..., : self.min_tokens_to_keep] = 0
+        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+
+        scores = scores.masked_fill(indices_to_remove, self.filter_value)
+        return scores

TTS/tts/layers/tortoise/audio_utils.py

@@ -0,0 +1,177 @@
+import os
+from glob import glob
+from typing import Dict, List
+
+import librosa
+import numpy as np
+import torch
+import torchaudio
+from scipy.io.wavfile import read
+
+from TTS.utils.audio.torch_transforms import TorchSTFT
+
+
+def load_wav_to_torch(full_path):
+    sampling_rate, data = read(full_path)
+    if data.dtype == np.int32:
+        norm_fix = 2**31
+    elif data.dtype == np.int16:
+        norm_fix = 2**15
+    elif data.dtype == np.float16 or data.dtype == np.float32:
+        norm_fix = 1.0
+    else:
+        raise NotImplementedError(f"Provided data dtype not supported: {data.dtype}")
+    return (torch.FloatTensor(data.astype(np.float32)) / norm_fix, sampling_rate)
+
+
+def check_audio(audio, audiopath: str):
+    # 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.
+    # '2' is arbitrarily chosen since it seems like audio will often "overdrive" the [-1,1] bounds.
+    if torch.any(audio > 2) or not torch.any(audio < 0):
+        print(f"Error with {audiopath}. Max={audio.max()} min={audio.min()}")
+    audio.clip_(-1, 1)
+
+
+def read_audio_file(audiopath: str):
+    if audiopath[-4:] == ".wav":
+        audio, lsr = load_wav_to_torch(audiopath)
+    elif audiopath[-4:] == ".mp3":
+        audio, lsr = librosa.load(audiopath, sr=None)
+        audio = torch.FloatTensor(audio)
+    else:
+        assert False, f"Unsupported audio format provided: {audiopath[-4:]}"
+
+    # Remove any channel data.
+    if len(audio.shape) > 1:
+        if audio.shape[0] < 5:
+            audio = audio[0]
+        else:
+            assert audio.shape[1] < 5
+            audio = audio[:, 0]
+
+    return audio, lsr
+
+
+def load_required_audio(audiopath: str):
+    audio, lsr = read_audio_file(audiopath)
+
+    audios = [torchaudio.functional.resample(audio, lsr, sampling_rate) for sampling_rate in (22050, 24000)]
+    for audio in audios:
+        check_audio(audio, audiopath)
+
+    return [audio.unsqueeze(0) for audio in audios]
+
+
+def load_audio(audiopath, sampling_rate):
+    audio, lsr = read_audio_file(audiopath)
+
+    if lsr != sampling_rate:
+        audio = torchaudio.functional.resample(audio, lsr, sampling_rate)
+    check_audio(audio, audiopath)
+
+    return audio.unsqueeze(0)
+
+
+TACOTRON_MEL_MAX = 2.3143386840820312
+TACOTRON_MEL_MIN = -11.512925148010254
+
+
+def denormalize_tacotron_mel(norm_mel):
+    return ((norm_mel + 1) / 2) * (TACOTRON_MEL_MAX - TACOTRON_MEL_MIN) + TACOTRON_MEL_MIN
+
+
+def normalize_tacotron_mel(mel):
+    return 2 * ((mel - TACOTRON_MEL_MIN) / (TACOTRON_MEL_MAX - TACOTRON_MEL_MIN)) - 1
+
+
+def dynamic_range_compression(x, C=1, clip_val=1e-5):
+    """
+    PARAMS
+    ------
+    C: compression factor
+    """
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def dynamic_range_decompression(x, C=1):
+    """
+    PARAMS
+    ------
+    C: compression factor used to compress
+    """
+    return torch.exp(x) / C
+
+
+def get_voices(extra_voice_dirs: List[str] = []):
+    dirs = extra_voice_dirs
+    voices: Dict[str, List[str]] = {}
+    for d in dirs:
+        subs = os.listdir(d)
+        for sub in subs:
+            subj = os.path.join(d, sub)
+            if os.path.isdir(subj):
+                voices[sub] = list(glob(f"{subj}/*.wav")) + list(glob(f"{subj}/*.mp3")) + list(glob(f"{subj}/*.pth"))
+    return voices
+
+
+def load_voice(voice: str, extra_voice_dirs: List[str] = []):
+    if voice == "random":
+        return None, None
+
+    voices = get_voices(extra_voice_dirs)
+    paths = voices[voice]
+    if len(paths) == 1 and paths[0].endswith(".pth"):
+        return None, torch.load(paths[0])
+    else:
+        conds = []
+        for cond_path in paths:
+            c = load_required_audio(cond_path)
+            conds.append(c)
+        return conds, None
+
+
+def load_voices(voices: List[str], extra_voice_dirs: List[str] = []):
+    latents = []
+    clips = []
+    for voice in voices:
+        if voice == "random":
+            if len(voices) > 1:
+                print("Cannot combine a random voice with a non-random voice. Just using a random voice.")
+            return None, None
+        clip, latent = load_voice(voice, extra_voice_dirs)
+        if latent is None:
+            assert (
+                len(latents) == 0
+            ), "Can only combine raw audio voices or latent voices, not both. Do it yourself if you want this."
+            clips.extend(clip)
+        elif clip is None:
+            assert (
+                len(clips) == 0
+            ), "Can only combine raw audio voices or latent voices, not both. Do it yourself if you want this."
+            latents.append(latent)
+    if len(latents) == 0:
+        return clips, None
+    else:
+        latents_0 = torch.stack([l[0] for l in latents], dim=0).mean(dim=0)
+        latents_1 = torch.stack([l[1] for l in latents], dim=0).mean(dim=0)
+        latents = (latents_0, latents_1)
+        return None, latents
+
+
+def wav_to_univnet_mel(wav, do_normalization=False, device="cuda"):
+    stft = TorchSTFT(
+        n_fft=1024,
+        hop_length=256,
+        win_length=1024,
+        use_mel=True,
+        n_mels=100,
+        sample_rate=24000,
+        mel_fmin=0,
+        mel_fmax=12000,
+    )
+    stft = stft.to(device)
+    mel = stft(wav)
+    mel = dynamic_range_compression(mel)
+    if do_normalization:
+        mel = normalize_tacotron_mel(mel)
+    return mel

TTS/tts/layers/tortoise/autoregressive.py

@@ -0,0 +1,631 @@
+# AGPL: a notification must be added stating that changes have been made to that file.
+import functools
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import GPT2Config, GPT2PreTrainedModel, LogitsProcessorList
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+
+from TTS.tts.layers.tortoise.arch_utils import AttentionBlock, TypicalLogitsWarper
+
+
+def null_position_embeddings(range, dim):
+    return torch.zeros((range.shape[0], range.shape[1], dim), device=range.device)
+
+
+def _p(t):
+    return t and (len(t), len(t[0]), t[0][0].shape)  # kv_cache debug
+
+
+class ResBlock(nn.Module):
+    """
+    Basic residual convolutional block that uses GroupNorm.
+    """
+
+    def __init__(self, chan):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv1d(chan, chan, kernel_size=3, padding=1),
+            nn.GroupNorm(chan // 8, chan),
+            nn.ReLU(),
+            nn.Conv1d(chan, chan, kernel_size=3, padding=1),
+            nn.GroupNorm(chan // 8, chan),
+        )
+
+    def forward(self, x):
+        return F.relu(self.net(x) + x)
+
+
+class GPT2InferenceModel(GPT2PreTrainedModel):
+    def __init__(self, config, gpt, text_pos_emb, embeddings, norm, linear, kv_cache):
+        super().__init__(config)
+        self.transformer = gpt
+        self.text_pos_embedding = text_pos_emb
+        self.embeddings = embeddings
+        self.lm_head = nn.Sequential(norm, linear)
+        self.kv_cache = kv_cache
+
+    def store_mel_emb(self, mel_emb):
+        self.cached_mel_emb = mel_emb
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        token_type_ids = kwargs.get("token_type_ids", None)  # usually None
+        if not self.kv_cache:
+            past_key_values = None
+        # only last token for inputs_ids if past is defined in kwargs
+        if past_key_values:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
+
+        attention_mask = kwargs.get("attention_mask", None)
+        position_ids = kwargs.get("position_ids", None)
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+        else:
+            position_ids = None
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "use_cache": kwargs.get("use_cache"),
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
+            "token_type_ids": token_type_ids,
+        }
+
+    def forward(
+        self,
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        assert self.cached_mel_emb is not None
+        assert inputs_embeds is None  # Not supported by this inference model.
+        assert labels is None  # Training not supported by this inference model.
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Create embedding
+        mel_len = self.cached_mel_emb.shape[1]
+        if input_ids.shape[1] != 1:
+            text_inputs = input_ids[:, mel_len:]
+            text_emb = self.embeddings(text_inputs)
+            text_emb = text_emb + self.text_pos_embedding(text_emb)
+            if self.cached_mel_emb.shape[0] != text_emb.shape[0]:
+                mel_emb = self.cached_mel_emb.repeat_interleave(text_emb.shape[0] // self.cached_mel_emb.shape[0], 0)
+            else:  # this outcome only occurs once per loop in most cases
+                mel_emb = self.cached_mel_emb
+            emb = torch.cat([mel_emb, text_emb], dim=1)
+        else:
+            emb = self.embeddings(input_ids)
+            emb = emb + self.text_pos_embedding.get_fixed_embedding(
+                attention_mask.shape[1] - mel_len, attention_mask.device
+            )
+
+        transformer_outputs = self.transformer(
+            inputs_embeds=emb,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+        lm_logits = self.lm_head(hidden_states)
+
+        if not return_dict:
+            return (lm_logits,) + transformer_outputs[1:]
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=None,
+            logits=lm_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+            cross_attentions=transformer_outputs.cross_attentions,
+        )
+
+    @staticmethod
+    def _reorder_cache(past, beam_idx):
+        """
+        This function is used to re-order the :obj:`past_key_values` cache if
+        :meth:`~transformers.PreTrainedModel.beam_search` or :meth:`~transformers.PreTrainedModel.beam_sample` is
+        called. This is required to match :obj:`past_key_values` with the correct beam_idx at every generation step.
+        """
+        return tuple(
+            tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
+            for layer_past in past
+        )
+
+
+class ConditioningEncoder(nn.Module):
+    def __init__(
+        self,
+        spec_dim,
+        embedding_dim,
+        attn_blocks=6,
+        num_attn_heads=4,
+        do_checkpointing=False,
+        mean=False,
+    ):
+        super().__init__()
+        attn = []
+        self.init = nn.Conv1d(spec_dim, embedding_dim, kernel_size=1)
+        for a in range(attn_blocks):
+            attn.append(AttentionBlock(embedding_dim, num_attn_heads))
+        self.attn = nn.Sequential(*attn)
+        self.dim = embedding_dim
+        self.do_checkpointing = do_checkpointing
+        self.mean = mean
+
+    def forward(self, x):
+        h = self.init(x)
+        h = self.attn(h)
+        if self.mean:
+            return h.mean(dim=2)
+        else:
+            return h[:, :, 0]
+
+
+class LearnedPositionEmbeddings(nn.Module):
+    def __init__(self, seq_len, model_dim, init=0.02):
+        super().__init__()
+        self.emb = nn.Embedding(seq_len, model_dim)
+        # Initializing this way is standard for GPT-2
+        self.emb.weight.data.normal_(mean=0.0, std=init)
+
+    def forward(self, x):
+        sl = x.shape[1]
+        return self.emb(torch.arange(0, sl, device=x.device))
+
+    def get_fixed_embedding(self, ind, dev):
+        return self.emb(torch.arange(0, ind, device=dev))[ind - 1 : ind]
+
+
+def build_hf_gpt_transformer(layers, model_dim, heads, max_mel_seq_len, max_text_seq_len, checkpointing):
+    """
+    GPT-2 implemented by the HuggingFace library.
+    """
+    from transformers import GPT2Config, GPT2Model
+
+    gpt_config = GPT2Config(
+        vocab_size=256,  # Unused.
+        n_positions=max_mel_seq_len + max_text_seq_len,
+        n_ctx=max_mel_seq_len + max_text_seq_len,
+        n_embd=model_dim,
+        n_layer=layers,
+        n_head=heads,
+        gradient_checkpointing=checkpointing,
+        use_cache=not checkpointing,
+    )
+    gpt = GPT2Model(gpt_config)
+    # Override the built in positional embeddings
+    del gpt.wpe  # TODO: figure out relevance in fixing exported model definition: Embedding(1012, 1024)
+    gpt.wpe = functools.partial(null_position_embeddings, dim=model_dim)
+    # Built-in token embeddings are unused.
+    del gpt.wte
+    return (
+        gpt,
+        LearnedPositionEmbeddings(max_mel_seq_len, model_dim),
+        LearnedPositionEmbeddings(max_text_seq_len, model_dim),
+        None,
+        None,
+    )
+
+
+class MelEncoder(nn.Module):
+    def __init__(self, channels, mel_channels=80, resblocks_per_reduction=2):
+        super().__init__()
+        self.channels = channels
+        self.encoder = nn.Sequential(
+            nn.Conv1d(mel_channels, channels // 4, kernel_size=3, padding=1),
+            nn.Sequential(*[ResBlock(channels // 4) for _ in range(resblocks_per_reduction)]),
+            nn.Conv1d(channels // 4, channels // 2, kernel_size=3, stride=2, padding=1),
+            nn.GroupNorm(channels // 16, channels // 2),
+            nn.ReLU(),
+            nn.Sequential(*[ResBlock(channels // 2) for _ in range(resblocks_per_reduction)]),
+            nn.Conv1d(channels // 2, channels, kernel_size=3, stride=2, padding=1),
+            nn.GroupNorm(channels // 8, channels),
+            nn.ReLU(),
+            nn.Sequential(*[ResBlock(channels) for _ in range(resblocks_per_reduction)]),
+        )
+        self.reduction = 4
+
+    def forward(self, x):
+        for e in self.encoder:
+            x = e(x)
+        return x.permute(0, 2, 1)
+
+
+class UnifiedVoice(nn.Module):
+    def __init__(
+        self,
+        layers=8,
+        model_dim=512,
+        heads=8,
+        max_text_tokens=120,
+        max_mel_tokens=250,
+        max_conditioning_inputs=1,
+        mel_length_compression=1024,
+        number_text_tokens=256,
+        start_text_token=None,
+        number_mel_codes=8194,
+        start_mel_token=8192,
+        stop_mel_token=8193,
+        train_solo_embeddings=False,
+        use_mel_codes_as_input=True,
+        checkpointing=True,
+        types=1,
+    ):
+        """
+        Args:
+            layers: Number of layers in transformer stack.
+            model_dim: Operating dimensions of the transformer
+            heads: Number of transformer heads. Must be divisible by model_dim. Recommend model_dim//64
+            max_text_tokens: Maximum number of text tokens that will be encountered by model.
+            max_mel_tokens: Maximum number of MEL tokens that will be encountered by model.
+            max_conditioning_inputs: Maximum number of conditioning inputs provided to the model. If (1), conditioning input can be of format (b,80,s), otherwise (b,n,80,s).
+            mel_length_compression: The factor between <number_input_samples> and <mel_tokens>. Used to compute MEL code padding given wav input length.
+            number_text_tokens:
+            start_text_token:
+            stop_text_token:
+            number_mel_codes:
+            start_mel_token:
+            stop_mel_token:
+            train_solo_embeddings:
+            use_mel_codes_as_input:
+            checkpointing:
+        """
+        super().__init__()
+
+        self.number_text_tokens = number_text_tokens
+        self.start_text_token = number_text_tokens * types if start_text_token is None else start_text_token
+        self.stop_text_token = 0
+        self.number_mel_codes = number_mel_codes
+        self.start_mel_token = start_mel_token
+        self.stop_mel_token = stop_mel_token
+        self.layers = layers
+        self.heads = heads
+        self.max_mel_tokens = max_mel_tokens
+        self.max_text_tokens = max_text_tokens
+        self.model_dim = model_dim
+        self.max_conditioning_inputs = max_conditioning_inputs
+        self.mel_length_compression = mel_length_compression
+        self.conditioning_encoder = ConditioningEncoder(80, model_dim, num_attn_heads=heads)
+        self.text_embedding = nn.Embedding(self.number_text_tokens * types + 1, model_dim)
+        if use_mel_codes_as_input:
+            self.mel_embedding = nn.Embedding(self.number_mel_codes, model_dim)
+        else:
+            self.mel_embedding = MelEncoder(model_dim, resblocks_per_reduction=1)
+        (
+            self.gpt,
+            self.mel_pos_embedding,
+            self.text_pos_embedding,
+            self.mel_layer_pos_embedding,
+            self.text_layer_pos_embedding,
+        ) = build_hf_gpt_transformer(
+            layers,
+            model_dim,
+            heads,
+            self.max_mel_tokens + 2 + self.max_conditioning_inputs,
+            self.max_text_tokens + 2,
+            checkpointing,
+        )
+        if train_solo_embeddings:
+            self.mel_solo_embedding = nn.Parameter(torch.randn(1, 1, model_dim) * 0.02, requires_grad=True)
+            self.text_solo_embedding = nn.Parameter(torch.randn(1, 1, model_dim) * 0.02, requires_grad=True)
+        else:
+            self.mel_solo_embedding = 0
+            self.text_solo_embedding = 0
+
+        self.final_norm = nn.LayerNorm(model_dim)
+        self.text_head = nn.Linear(model_dim, self.number_text_tokens * types + 1)
+        self.mel_head = nn.Linear(model_dim, self.number_mel_codes)
+
+        # Initialize the embeddings per the GPT-2 scheme
+        embeddings = [self.text_embedding]
+        if use_mel_codes_as_input:
+            embeddings.append(self.mel_embedding)
+        for module in embeddings:
+            module.weight.data.normal_(mean=0.0, std=0.02)
+
+    def post_init_gpt2_config(self, kv_cache=True):
+        seq_length = self.max_mel_tokens + self.max_text_tokens + 2
+        gpt_config = GPT2Config(
+            vocab_size=self.max_mel_tokens,
+            n_positions=seq_length,
+            n_ctx=seq_length,
+            n_embd=self.model_dim,
+            n_layer=self.layers,
+            n_head=self.heads,
+            gradient_checkpointing=False,
+            use_cache=True,
+        )
+        self.inference_model = GPT2InferenceModel(
+            gpt_config,
+            self.gpt,
+            self.mel_pos_embedding,
+            self.mel_embedding,
+            self.final_norm,
+            self.mel_head,
+            kv_cache=kv_cache,
+        )
+        # self.inference_model = PrunedGPT2InferenceModel(gpt_config, self.gpt, self.mel_pos_embedding, self.mel_embedding, self.final_norm, self.mel_head)
+        self.gpt.wte = self.mel_embedding
+        # self.inference_model.save_pretrained("")
+
+    def build_aligned_inputs_and_targets(self, input, start_token, stop_token):
+        inp = F.pad(input, (1, 0), value=start_token)
+        tar = F.pad(input, (0, 1), value=stop_token)
+        return inp, tar
+
+    def set_mel_padding(self, mel_input_tokens, wav_lengths):
+        """
+        Given mel tokens that are derived from a padded audio clip and the actual lengths of each batch element in
+        that audio clip, reformats the tokens with STOP_MEL_TOKEN in place of the zero padding. This is required
+        preformatting to create a working TTS model.
+        """
+        # Set padding areas within MEL (currently it is coded with the MEL code for <zero>).
+        mel_lengths = torch.div(wav_lengths, self.mel_length_compression, rounding_mode="trunc")
+        for b in range(len(mel_lengths)):
+            actual_end = (
+                mel_lengths[b] + 1
+            )  # Due to the convolutional nature of how these tokens are generated, it would be best if the model predicts a token past the actual last token.
+            if actual_end < mel_input_tokens.shape[-1]:
+                mel_input_tokens[b, actual_end:] = self.stop_mel_token
+        return mel_input_tokens
+
+    def get_logits(
+        self,
+        speech_conditioning_inputs,
+        first_inputs,
+        first_head,
+        second_inputs=None,
+        second_head=None,
+        get_attns=False,
+        return_latent=False,
+    ):
+        if second_inputs is not None:
+            emb = torch.cat([speech_conditioning_inputs, first_inputs, second_inputs], dim=1)
+        else:
+            emb = torch.cat([speech_conditioning_inputs, first_inputs], dim=1)
+
+        gpt_out = self.gpt(inputs_embeds=emb, return_dict=True, output_attentions=get_attns)
+        if get_attns:
+            return gpt_out.attentions
+
+        enc = gpt_out.last_hidden_state[:, 1:]  # The first logit is tied to the speech_conditioning_input
+        enc = self.final_norm(enc)
+
+        if return_latent:
+            return (
+                enc[
+                    :,
+                    speech_conditioning_inputs.shape[1] : speech_conditioning_inputs.shape[1] + first_inputs.shape[1],
+                ],
+                enc[:, -second_inputs.shape[1] :],
+            )
+
+        first_logits = enc[:, : first_inputs.shape[1]]
+        first_logits = first_head(first_logits)
+        first_logits = first_logits.permute(0, 2, 1)
+        if second_inputs is not None:
+            second_logits = enc[:, -second_inputs.shape[1] :]
+            second_logits = second_head(second_logits)
+            second_logits = second_logits.permute(0, 2, 1)
+            return first_logits, second_logits
+        else:
+            return first_logits
+
+    def get_conditioning(self, speech_conditioning_input):
+        speech_conditioning_input = (
+            speech_conditioning_input.unsqueeze(1)
+            if len(speech_conditioning_input.shape) == 3
+            else speech_conditioning_input
+        )
+        conds = []
+        for j in range(speech_conditioning_input.shape[1]):
+            conds.append(self.conditioning_encoder(speech_conditioning_input[:, j]))
+        conds = torch.stack(conds, dim=1)
+        conds = conds.mean(dim=1)
+        return conds
+
+    def forward(
+        self,
+        speech_conditioning_latent,
+        text_inputs,
+        text_lengths,
+        mel_codes,
+        wav_lengths,
+        types=None,
+        text_first=True,
+        raw_mels=None,
+        return_attentions=False,
+        return_latent=False,
+        clip_inputs=True,
+    ):
+        """
+        Forward pass that uses both text and voice in either text conditioning mode or voice conditioning mode
+        (actuated by `text_first`).
+
+        speech_conditioning_input: MEL float tensor, (b,1024)
+        text_inputs: long tensor, (b,t)
+        text_lengths: long tensor, (b,)
+        mel_inputs:  long tensor, (b,m)
+        wav_lengths: long tensor, (b,)
+        raw_mels: MEL float tensor (b,80,s)
+
+        If return_attentions is specified, only logits are returned.
+        If return_latent is specified, loss & logits are not computed or returned. Only the predicted latents are returned.
+        If clip_inputs is True, the inputs will be clipped to the smallest input size across each input modality.
+        """
+        # Types are expressed by expanding the text embedding space.
+        if types is not None:
+            text_inputs = text_inputs * (1 + types).unsqueeze(-1)
+
+        if clip_inputs:
+            # This model will receive micro-batches with a ton of padding for both the text and MELs. Ameliorate this by
+            # chopping the inputs by the maximum actual length.
+            max_text_len = text_lengths.max()
+            text_inputs = text_inputs[:, :max_text_len]
+            max_mel_len = wav_lengths.max() // self.mel_length_compression
+            mel_codes = mel_codes[:, :max_mel_len]
+            if raw_mels is not None:
+                raw_mels = raw_mels[:, :, : max_mel_len * 4]
+        mel_codes = self.set_mel_padding(mel_codes, wav_lengths)
+        text_inputs = F.pad(text_inputs, (0, 1), value=self.stop_text_token)
+        mel_codes = F.pad(mel_codes, (0, 1), value=self.stop_mel_token)
+
+        conds = speech_conditioning_latent.unsqueeze(1)
+        text_inputs, text_targets = self.build_aligned_inputs_and_targets(
+            text_inputs, self.start_text_token, self.stop_text_token
+        )
+        text_emb = self.text_embedding(text_inputs) + self.text_pos_embedding(text_inputs)
+        mel_codes, mel_targets = self.build_aligned_inputs_and_targets(
+            mel_codes, self.start_mel_token, self.stop_mel_token
+        )
+        if raw_mels is not None:
+            mel_inp = F.pad(raw_mels, (0, 8))
+        else:
+            mel_inp = mel_codes
+        mel_emb = self.mel_embedding(mel_inp)
+        mel_emb = mel_emb + self.mel_pos_embedding(mel_codes)
+
+        if text_first:
+            text_logits, mel_logits = self.get_logits(
+                conds,
+                text_emb,
+                self.text_head,
+                mel_emb,
+                self.mel_head,
+                get_attns=return_attentions,
+                return_latent=return_latent,
+            )
+            if return_latent:
+                return mel_logits[
+                    :, :-2
+                ]  # Despite the name, these are not logits. Strip off the two tokens added by this forward pass.
+        else:
+            mel_logits, text_logits = self.get_logits(
+                conds,
+                mel_emb,
+                self.mel_head,
+                text_emb,
+                self.text_head,
+                get_attns=return_attentions,
+                return_latent=return_latent,
+            )
+            if return_latent:
+                return text_logits[
+                    :, :-2
+                ]  # Despite the name, these are not logits. Strip off the two tokens added by this forward pass.
+
+        if return_attentions:
+            return mel_logits
+        loss_text = F.cross_entropy(text_logits, text_targets.long())
+        loss_mel = F.cross_entropy(mel_logits, mel_targets.long())
+        return loss_text.mean(), loss_mel.mean(), mel_logits
+
+    def inference_speech(
+        self,
+        speech_conditioning_latent,
+        text_inputs,
+        input_tokens=None,
+        num_return_sequences=1,
+        max_generate_length=None,
+        typical_sampling=False,
+        typical_mass=0.9,
+        **hf_generate_kwargs,
+    ):
+        text_inputs = F.pad(text_inputs, (0, 1), value=self.stop_text_token)
+        text_inputs, text_targets = self.build_aligned_inputs_and_targets(
+            text_inputs, self.start_text_token, self.stop_text_token
+        )
+        text_emb = self.text_embedding(text_inputs) + self.text_pos_embedding(text_inputs)
+
+        conds = speech_conditioning_latent.unsqueeze(1)
+        emb = torch.cat([conds, text_emb], dim=1)
+        self.inference_model.store_mel_emb(emb)
+
+        fake_inputs = torch.full(
+            (
+                emb.shape[0],
+                conds.shape[1] + emb.shape[1],
+            ),
+            fill_value=1,
+            dtype=torch.long,
+            device=text_inputs.device,
+        )
+        fake_inputs[:, -1] = self.start_mel_token
+        trunc_index = fake_inputs.shape[1]
+        if input_tokens is None:
+            inputs = fake_inputs
+        else:
+            assert (
+                num_return_sequences % input_tokens.shape[0] == 0
+            ), "The number of return sequences must be divisible by the number of input sequences"
+            fake_inputs = fake_inputs.repeat(num_return_sequences, 1)
+            input_tokens = input_tokens.repeat(num_return_sequences // input_tokens.shape[0], 1)
+            inputs = torch.cat([fake_inputs, input_tokens], dim=1)
+
+        logits_processor = (
+            LogitsProcessorList([TypicalLogitsWarper(mass=typical_mass)]) if typical_sampling else LogitsProcessorList()
+        )  # TODO disable this
+        max_length = (
+            trunc_index + self.max_mel_tokens - 1 if max_generate_length is None else trunc_index + max_generate_length
+        )
+        gen = self.inference_model.generate(
+            inputs,
+            bos_token_id=self.start_mel_token,
+            pad_token_id=self.stop_mel_token,
+            eos_token_id=self.stop_mel_token,
+            max_length=max_length,
+            logits_processor=logits_processor,
+            num_return_sequences=num_return_sequences,
+            **hf_generate_kwargs,
+        )
+        return gen[:, trunc_index:]
+
+
+if __name__ == "__main__":
+    gpt = UnifiedVoice(
+        model_dim=256,
+        heads=4,
+        train_solo_embeddings=True,
+        use_mel_codes_as_input=True,
+        max_conditioning_inputs=4,
+    )
+    l = gpt(
+        torch.randn(2, 3, 80, 800),
+        torch.randint(high=120, size=(2, 120)),
+        torch.tensor([32, 120]),
+        torch.randint(high=8192, size=(2, 250)),
+        torch.tensor([250 * 256, 195 * 256]),
+    )
+    gpt.text_forward(
+        torch.randn(2, 80, 800),
+        torch.randint(high=50, size=(2, 80)),
+        torch.tensor([32, 80]),
+    )

TTS/tts/layers/tortoise/classifier.py

@@ -0,0 +1,144 @@
+import torch
+import torch.nn as nn
+
+from TTS.tts.layers.tortoise.arch_utils import AttentionBlock, Downsample, Upsample, normalization, zero_module
+
+
+class ResBlock(nn.Module):
+    def __init__(
+        self,
+        channels,
+        dropout,
+        out_channels=None,
+        use_conv=False,
+        use_scale_shift_norm=False,
+        dims=2,
+        up=False,
+        down=False,
+        kernel_size=3,
+        do_checkpoint=True,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_scale_shift_norm = use_scale_shift_norm
+        self.do_checkpoint = do_checkpoint
+        padding = 1 if kernel_size == 3 else 2
+
+        self.in_layers = nn.Sequential(
+            normalization(channels),
+            nn.SiLU(),
+            nn.Conv1d(channels, self.out_channels, kernel_size, padding=padding),
+        )
+
+        self.updown = up or down
+
+        if up:
+            self.h_upd = Upsample(channels, False, dims)
+            self.x_upd = Upsample(channels, False, dims)
+        elif down:
+            self.h_upd = Downsample(channels, False, dims)
+            self.x_upd = Downsample(channels, False, dims)
+        else:
+            self.h_upd = self.x_upd = nn.Identity()
+
+        self.out_layers = nn.Sequential(
+            normalization(self.out_channels),
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            zero_module(nn.Conv1d(self.out_channels, self.out_channels, kernel_size, padding=padding)),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        elif use_conv:
+            self.skip_connection = nn.Conv1d(dims, channels, self.out_channels, kernel_size, padding=padding)
+        else:
+            self.skip_connection = nn.Conv1d(dims, channels, self.out_channels, 1)
+
+    def forward(self, x):
+        if self.updown:
+            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+            h = in_rest(x)
+            h = self.h_upd(h)
+            x = self.x_upd(x)
+            h = in_conv(h)
+        else:
+            h = self.in_layers(x)
+        h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+
+class AudioMiniEncoder(nn.Module):
+    def __init__(
+        self,
+        spec_dim,
+        embedding_dim,
+        base_channels=128,
+        depth=2,
+        resnet_blocks=2,
+        attn_blocks=4,
+        num_attn_heads=4,
+        dropout=0,
+        downsample_factor=2,
+        kernel_size=3,
+    ):
+        super().__init__()
+        self.init = nn.Sequential(nn.Conv1d(spec_dim, base_channels, 3, padding=1))
+        ch = base_channels
+        res = []
+        self.layers = depth
+        for l in range(depth):
+            for r in range(resnet_blocks):
+                res.append(ResBlock(ch, dropout, do_checkpoint=False, kernel_size=kernel_size))
+            res.append(Downsample(ch, use_conv=True, out_channels=ch * 2, factor=downsample_factor))
+            ch *= 2
+        self.res = nn.Sequential(*res)
+        self.final = nn.Sequential(normalization(ch), nn.SiLU(), nn.Conv1d(ch, embedding_dim, 1))
+        attn = []
+        for a in range(attn_blocks):
+            attn.append(AttentionBlock(embedding_dim, num_attn_heads, do_checkpoint=False))
+        self.attn = nn.Sequential(*attn)
+        self.dim = embedding_dim
+
+    def forward(self, x):
+        h = self.init(x)
+        h = self.res(h)
+        h = self.final(h)
+        for blk in self.attn:
+            h = blk(h)
+        return h[:, :, 0]
+
+
+class AudioMiniEncoderWithClassifierHead(nn.Module):
+    def __init__(self, classes, distribute_zero_label=True, **kwargs):
+        super().__init__()
+        self.enc = AudioMiniEncoder(**kwargs)
+        self.head = nn.Linear(self.enc.dim, classes)
+        self.num_classes = classes
+        self.distribute_zero_label = distribute_zero_label
+
+    def forward(self, x, labels=None):
+        h = self.enc(x)
+        logits = self.head(h)
+        if labels is None:
+            return logits
+        else:
+            if self.distribute_zero_label:
+                oh_labels = nn.functional.one_hot(labels, num_classes=self.num_classes)
+                zeros_indices = (labels == 0).unsqueeze(-1)
+                # Distribute 20% of the probability mass on all classes when zero is specified, to compensate for dataset noise.
+                zero_extra_mass = torch.full_like(
+                    oh_labels,
+                    dtype=torch.float,
+                    fill_value=0.2 / (self.num_classes - 1),
+                )
+                zero_extra_mass[:, 0] = -0.2
+                zero_extra_mass = zero_extra_mass * zeros_indices
+                oh_labels = oh_labels + zero_extra_mass
+            else:
+                oh_labels = labels
+            loss = nn.functional.cross_entropy(logits, oh_labels)
+            return loss

TTS/tts/layers/tortoise/clvp.py

@@ -0,0 +1,159 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import einsum
+
+from TTS.tts.layers.tortoise.arch_utils import CheckpointedXTransformerEncoder
+from TTS.tts.layers.tortoise.transformer import Transformer
+from TTS.tts.layers.tortoise.xtransformers import Encoder
+
+
+def exists(val):
+    return val is not None
+
+
+def masked_mean(t, mask, dim=1):
+    t = t.masked_fill(~mask[:, :, None], 0.0)
+    return t.sum(dim=1) / mask.sum(dim=1)[..., None]
+
+
+class CLVP(nn.Module):
+    """
+    CLIP model retrofitted for performing contrastive evaluation between tokenized audio data and the corresponding
+    transcribed text.
+
+    Originally from https://github.com/lucidrains/DALLE-pytorch/blob/main/dalle_pytorch/dalle_pytorch.py
+    """
+
+    def __init__(
+        self,
+        *,
+        dim_text=512,
+        dim_speech=512,
+        dim_latent=512,
+        num_text_tokens=256,
+        text_enc_depth=6,
+        text_seq_len=120,
+        text_heads=8,
+        num_speech_tokens=8192,
+        speech_enc_depth=6,
+        speech_heads=8,
+        speech_seq_len=250,
+        text_mask_percentage=0,
+        voice_mask_percentage=0,
+        wav_token_compression=1024,
+        use_xformers=False,
+    ):
+        super().__init__()
+        self.text_emb = nn.Embedding(num_text_tokens, dim_text)
+        self.to_text_latent = nn.Linear(dim_text, dim_latent, bias=False)
+
+        self.speech_emb = nn.Embedding(num_speech_tokens, dim_speech)
+        self.to_speech_latent = nn.Linear(dim_speech, dim_latent, bias=False)
+
+        if use_xformers:
+            self.text_transformer = CheckpointedXTransformerEncoder(
+                needs_permute=False,
+                exit_permute=False,
+                max_seq_len=-1,
+                attn_layers=Encoder(
+                    dim=dim_text,
+                    depth=text_enc_depth,
+                    heads=text_heads,
+                    ff_dropout=0.1,
+                    ff_mult=2,
+                    attn_dropout=0.1,
+                    use_rmsnorm=True,
+                    ff_glu=True,
+                    rotary_pos_emb=True,
+                ),
+            )
+            self.speech_transformer = CheckpointedXTransformerEncoder(
+                needs_permute=False,
+                exit_permute=False,
+                max_seq_len=-1,
+                attn_layers=Encoder(
+                    dim=dim_speech,
+                    depth=speech_enc_depth,
+                    heads=speech_heads,
+                    ff_dropout=0.1,
+                    ff_mult=2,
+                    attn_dropout=0.1,
+                    use_rmsnorm=True,
+                    ff_glu=True,
+                    rotary_pos_emb=True,
+                ),
+            )
+        else:
+            self.text_transformer = Transformer(
+                causal=False, seq_len=text_seq_len, dim=dim_text, depth=text_enc_depth, heads=text_heads
+            )
+            self.speech_transformer = Transformer(
+                causal=False, seq_len=speech_seq_len, dim=dim_speech, depth=speech_enc_depth, heads=speech_heads
+            )
+
+        self.temperature = nn.Parameter(torch.tensor(1.0))
+        self.text_mask_percentage = text_mask_percentage
+        self.voice_mask_percentage = voice_mask_percentage
+        self.wav_token_compression = wav_token_compression
+        self.xformers = use_xformers
+        if not use_xformers:
+            self.text_pos_emb = nn.Embedding(text_seq_len, dim_text)
+            self.speech_pos_emb = nn.Embedding(num_speech_tokens, dim_speech)
+
+    def forward(self, text, speech_tokens, return_loss=False):
+        b, device = text.shape[0], text.device
+        if self.training:
+            text_mask = torch.rand_like(text.float()) > self.text_mask_percentage
+            voice_mask = torch.rand_like(speech_tokens.float()) > self.voice_mask_percentage
+        else:
+            text_mask = torch.ones_like(text.float()).bool()
+            voice_mask = torch.ones_like(speech_tokens.float()).bool()
+
+        text_emb = self.text_emb(text)
+        speech_emb = self.speech_emb(speech_tokens)
+
+        if not self.xformers:
+            text_emb += self.text_pos_emb(torch.arange(text.shape[1], device=device))
+            speech_emb += self.speech_pos_emb(torch.arange(speech_emb.shape[1], device=device))
+
+        enc_text = self.text_transformer(text_emb, mask=text_mask)
+        enc_speech = self.speech_transformer(speech_emb, mask=voice_mask)
+
+        text_latents = masked_mean(enc_text, text_mask, dim=1)
+        speech_latents = masked_mean(enc_speech, voice_mask, dim=1)
+
+        text_latents = self.to_text_latent(text_latents)
+        speech_latents = self.to_speech_latent(speech_latents)
+
+        text_latents, speech_latents = map(lambda t: F.normalize(t, p=2, dim=-1), (text_latents, speech_latents))
+
+        temp = self.temperature.exp()
+
+        if not return_loss:
+            sim = einsum("n d, n d -> n", text_latents, speech_latents) * temp
+            return sim
+
+        sim = einsum("i d, j d -> i j", text_latents, speech_latents) * temp
+        labels = torch.arange(b, device=device)
+        loss = (F.cross_entropy(sim, labels) + F.cross_entropy(sim.t(), labels)) / 2
+        return loss
+
+
+if __name__ == "__main__":
+    clip = CLVP(text_mask_percentage=0.2, voice_mask_percentage=0.2)
+    clip(
+        torch.randint(0, 256, (2, 120)),
+        torch.tensor([50, 100]),
+        torch.randint(0, 8192, (2, 250)),
+        torch.tensor([101, 102]),
+        return_loss=True,
+    )
+    nonloss = clip(
+        torch.randint(0, 256, (2, 120)),
+        torch.tensor([50, 100]),
+        torch.randint(0, 8192, (2, 250)),
+        torch.tensor([101, 102]),
+        return_loss=False,
+    )
+    print(nonloss.shape)

TTS/tts/layers/tortoise/diffusion_decoder.py

@@ -0,0 +1,415 @@
+import math
+import random
+from abc import abstractmethod
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import autocast
+
+from TTS.tts.layers.tortoise.arch_utils import AttentionBlock, normalization
+
+
+def is_latent(t):
+    return t.dtype == torch.float
+
+
+def is_sequence(t):
+    return t.dtype == torch.long
+
+
+def timestep_embedding(timesteps, dim, max_period=10000):
+    """
+    Create sinusoidal timestep embeddings.
+
+    :param timesteps: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an [N x dim] Tensor of positional embeddings.
+    """
+    half = dim // 2
+    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
+        device=timesteps.device
+    )
+    args = timesteps[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    return embedding
+
+
+class TimestepBlock(nn.Module):
+    @abstractmethod
+    def forward(self, x, emb):
+        """
+        Apply the module to `x` given `emb` timestep embeddings.
+        """
+
+
+class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
+    def forward(self, x, emb):
+        for layer in self:
+            if isinstance(layer, TimestepBlock):
+                x = layer(x, emb)
+            else:
+                x = layer(x)
+        return x
+
+
+class ResBlock(TimestepBlock):
+    def __init__(
+        self,
+        channels,
+        emb_channels,
+        dropout,
+        out_channels=None,
+        dims=2,
+        kernel_size=3,
+        efficient_config=True,
+        use_scale_shift_norm=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.emb_channels = emb_channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_scale_shift_norm = use_scale_shift_norm
+        padding = {1: 0, 3: 1, 5: 2}[kernel_size]
+        eff_kernel = 1 if efficient_config else 3
+        eff_padding = 0 if efficient_config else 1
+
+        self.in_layers = nn.Sequential(
+            normalization(channels),
+            nn.SiLU(),
+            nn.Conv1d(channels, self.out_channels, eff_kernel, padding=eff_padding),
+        )
+
+        self.emb_layers = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(
+                emb_channels,
+                2 * self.out_channels if use_scale_shift_norm else self.out_channels,
+            ),
+        )
+        self.out_layers = nn.Sequential(
+            normalization(self.out_channels),
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            nn.Conv1d(self.out_channels, self.out_channels, kernel_size, padding=padding),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        else:
+            self.skip_connection = nn.Conv1d(channels, self.out_channels, eff_kernel, padding=eff_padding)
+
+    def forward(self, x, emb):
+        h = self.in_layers(x)
+        emb_out = self.emb_layers(emb).type(h.dtype)
+        while len(emb_out.shape) < len(h.shape):
+            emb_out = emb_out[..., None]
+        if self.use_scale_shift_norm:
+            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
+            scale, shift = torch.chunk(emb_out, 2, dim=1)
+            h = out_norm(h) * (1 + scale) + shift
+            h = out_rest(h)
+        else:
+            h = h + emb_out
+            h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+
+class DiffusionLayer(TimestepBlock):
+    def __init__(self, model_channels, dropout, num_heads):
+        super().__init__()
+        self.resblk = ResBlock(
+            model_channels,
+            model_channels,
+            dropout,
+            model_channels,
+            dims=1,
+            use_scale_shift_norm=True,
+        )
+        self.attn = AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True)
+
+    def forward(self, x, time_emb):
+        y = self.resblk(x, time_emb)
+        return self.attn(y)
+
+
+class DiffusionTts(nn.Module):
+    def __init__(
+        self,
+        model_channels=512,
+        num_layers=8,
+        in_channels=100,
+        in_latent_channels=512,
+        in_tokens=8193,
+        out_channels=200,  # mean and variance
+        dropout=0,
+        use_fp16=False,
+        num_heads=16,
+        # Parameters for regularization.
+        layer_drop=0.1,
+        unconditioned_percentage=0.1,  # This implements a mechanism similar to what is used in classifier-free training.
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.dropout = dropout
+        self.num_heads = num_heads
+        self.unconditioned_percentage = unconditioned_percentage
+        self.enable_fp16 = use_fp16
+        self.layer_drop = layer_drop
+
+        self.inp_block = nn.Conv1d(in_channels, model_channels, 3, 1, 1)
+        self.time_embed = nn.Sequential(
+            nn.Linear(model_channels, model_channels),
+            nn.SiLU(),
+            nn.Linear(model_channels, model_channels),
+        )
+
+        # Either code_converter or latent_converter is used, depending on what type of conditioning data is fed.
+        # This model is meant to be able to be trained on both for efficiency purposes - it is far less computationally
+        # complex to generate tokens, while generating latents will normally mean propagating through a deep autoregressive
+        # transformer network.
+        self.code_embedding = nn.Embedding(in_tokens, model_channels)
+        self.code_converter = nn.Sequential(
+            AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True),
+            AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True),
+            AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True),
+        )
+        self.code_norm = normalization(model_channels)
+        self.latent_conditioner = nn.Sequential(
+            nn.Conv1d(in_latent_channels, model_channels, 3, padding=1),
+            AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True),
+            AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True),
+            AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True),
+            AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True),
+        )
+        self.contextual_embedder = nn.Sequential(
+            nn.Conv1d(in_channels, model_channels, 3, padding=1, stride=2),
+            nn.Conv1d(model_channels, model_channels * 2, 3, padding=1, stride=2),
+            AttentionBlock(
+                model_channels * 2,
+                num_heads,
+                relative_pos_embeddings=True,
+                do_checkpoint=False,
+            ),
+            AttentionBlock(
+                model_channels * 2,
+                num_heads,
+                relative_pos_embeddings=True,
+                do_checkpoint=False,
+            ),
+            AttentionBlock(
+                model_channels * 2,
+                num_heads,
+                relative_pos_embeddings=True,
+                do_checkpoint=False,
+            ),
+            AttentionBlock(
+                model_channels * 2,
+                num_heads,
+                relative_pos_embeddings=True,
+                do_checkpoint=False,
+            ),
+            AttentionBlock(
+                model_channels * 2,
+                num_heads,
+                relative_pos_embeddings=True,
+                do_checkpoint=False,
+            ),
+        )
+        self.unconditioned_embedding = nn.Parameter(torch.randn(1, model_channels, 1))
+        self.conditioning_timestep_integrator = TimestepEmbedSequential(
+            DiffusionLayer(model_channels, dropout, num_heads),
+            DiffusionLayer(model_channels, dropout, num_heads),
+            DiffusionLayer(model_channels, dropout, num_heads),
+        )
+
+        self.integrating_conv = nn.Conv1d(model_channels * 2, model_channels, kernel_size=1)
+        self.mel_head = nn.Conv1d(model_channels, in_channels, kernel_size=3, padding=1)
+
+        self.layers = nn.ModuleList(
+            [DiffusionLayer(model_channels, dropout, num_heads) for _ in range(num_layers)]
+            + [
+                ResBlock(
+                    model_channels,
+                    model_channels,
+                    dropout,
+                    dims=1,
+                    use_scale_shift_norm=True,
+                )
+                for _ in range(3)
+            ]
+        )
+
+        self.out = nn.Sequential(
+            normalization(model_channels),
+            nn.SiLU(),
+            nn.Conv1d(model_channels, out_channels, 3, padding=1),
+        )
+
+    def get_grad_norm_parameter_groups(self):
+        groups = {
+            "minicoder": list(self.contextual_embedder.parameters()),
+            "layers": list(self.layers.parameters()),
+            "code_converters": list(self.code_embedding.parameters())
+            + list(self.code_converter.parameters())
+            + list(self.latent_conditioner.parameters())
+            + list(self.latent_conditioner.parameters()),
+            "timestep_integrator": list(self.conditioning_timestep_integrator.parameters())
+            + list(self.integrating_conv.parameters()),
+            "time_embed": list(self.time_embed.parameters()),
+        }
+        return groups
+
+    def get_conditioning(self, conditioning_input):
+        speech_conditioning_input = (
+            conditioning_input.unsqueeze(1) if len(conditioning_input.shape) == 3 else conditioning_input
+        )
+        conds = []
+        for j in range(speech_conditioning_input.shape[1]):
+            conds.append(self.contextual_embedder(speech_conditioning_input[:, j]))
+        conds = torch.cat(conds, dim=-1)
+        conds = conds.mean(dim=-1)
+        return conds
+
+    def timestep_independent(
+        self,
+        aligned_conditioning,
+        conditioning_latent,
+        expected_seq_len,
+        return_code_pred,
+    ):
+        # Shuffle aligned_latent to BxCxS format
+        if is_latent(aligned_conditioning):
+            aligned_conditioning = aligned_conditioning.permute(0, 2, 1)
+
+        cond_scale, cond_shift = torch.chunk(conditioning_latent, 2, dim=1)
+        if is_latent(aligned_conditioning):
+            code_emb = self.latent_conditioner(aligned_conditioning)
+        else:
+            code_emb = self.code_embedding(aligned_conditioning).permute(0, 2, 1)
+            code_emb = self.code_converter(code_emb)
+        code_emb = self.code_norm(code_emb) * (1 + cond_scale.unsqueeze(-1)) + cond_shift.unsqueeze(-1)
+
+        unconditioned_batches = torch.zeros((code_emb.shape[0], 1, 1), device=code_emb.device)
+        # Mask out the conditioning branch for whole batch elements, implementing something similar to classifier-free guidance.
+        if self.training and self.unconditioned_percentage > 0:
+            unconditioned_batches = (
+                torch.rand((code_emb.shape[0], 1, 1), device=code_emb.device) < self.unconditioned_percentage
+            )
+            code_emb = torch.where(
+                unconditioned_batches,
+                self.unconditioned_embedding.repeat(aligned_conditioning.shape[0], 1, 1),
+                code_emb,
+            )
+        expanded_code_emb = F.interpolate(code_emb, size=expected_seq_len, mode="nearest")
+
+        if not return_code_pred:
+            return expanded_code_emb
+        else:
+            mel_pred = self.mel_head(expanded_code_emb)
+            # Multiply mel_pred by !unconditioned_branches, which drops the gradient on unconditioned branches. This is because we don't want that gradient being used to train parameters through the codes_embedder as it unbalances contributions to that network from the MSE loss.
+            mel_pred = mel_pred * unconditioned_batches.logical_not()
+            return expanded_code_emb, mel_pred
+
+    def forward(
+        self,
+        x,
+        timesteps,
+        aligned_conditioning=None,
+        conditioning_latent=None,
+        precomputed_aligned_embeddings=None,
+        conditioning_free=False,
+        return_code_pred=False,
+    ):
+        """
+        Apply the model to an input batch.
+
+        :param x: an [N x C x ...] Tensor of inputs.
+        :param timesteps: a 1-D batch of timesteps.
+        :param aligned_conditioning: an aligned latent or sequence of tokens providing useful data about the sample to be produced.
+        :param conditioning_latent: a pre-computed conditioning latent; see get_conditioning().
+        :param precomputed_aligned_embeddings: Embeddings returned from self.timestep_independent()
+        :param conditioning_free: When set, all conditioning inputs (including tokens and conditioning_input) will not be considered.
+        :return: an [N x C x ...] Tensor of outputs.
+        """
+        assert precomputed_aligned_embeddings is not None or (
+            aligned_conditioning is not None and conditioning_latent is not None
+        )
+        assert not (
+            return_code_pred and precomputed_aligned_embeddings is not None
+        )  # These two are mutually exclusive.
+
+        unused_params = []
+        if conditioning_free:
+            code_emb = self.unconditioned_embedding.repeat(x.shape[0], 1, x.shape[-1])
+            unused_params.extend(list(self.code_converter.parameters()) + list(self.code_embedding.parameters()))
+            unused_params.extend(list(self.latent_conditioner.parameters()))
+        else:
+            if precomputed_aligned_embeddings is not None:
+                code_emb = precomputed_aligned_embeddings
+            else:
+                code_emb, mel_pred = self.timestep_independent(
+                    aligned_conditioning, conditioning_latent, x.shape[-1], True
+                )
+                if is_latent(aligned_conditioning):
+                    unused_params.extend(
+                        list(self.code_converter.parameters()) + list(self.code_embedding.parameters())
+                    )
+                else:
+                    unused_params.extend(list(self.latent_conditioner.parameters()))
+
+            unused_params.append(self.unconditioned_embedding)
+
+        time_emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
+        code_emb = self.conditioning_timestep_integrator(code_emb, time_emb)
+        x = self.inp_block(x)
+        x = torch.cat([x, code_emb], dim=1)
+        x = self.integrating_conv(x)
+        for i, lyr in enumerate(self.layers):
+            # Do layer drop where applicable. Do not drop first and last layers.
+            if (
+                self.training
+                and self.layer_drop > 0
+                and i != 0
+                and i != (len(self.layers) - 1)
+                and random.random() < self.layer_drop
+            ):
+                unused_params.extend(list(lyr.parameters()))
+            else:
+                # First and last blocks will have autocast disabled for improved precision.
+                with autocast(x.device.type, enabled=self.enable_fp16 and i != 0):
+                    x = lyr(x, time_emb)
+
+        x = x.float()
+        out = self.out(x)
+
+        # Involve probabilistic or possibly unused parameters in loss so we don't get DDP errors.
+        extraneous_addition = 0
+        for p in unused_params:
+            extraneous_addition = extraneous_addition + p.mean()
+        out = out + extraneous_addition * 0
+
+        if return_code_pred:
+            return out, mel_pred
+        return out
+
+
+if __name__ == "__main__":
+    clip = torch.randn(2, 100, 400)
+    aligned_latent = torch.randn(2, 388, 512)
+    aligned_sequence = torch.randint(0, 8192, (2, 100))
+    cond = torch.randn(2, 100, 400)
+    ts = torch.LongTensor([600, 600])
+    model = DiffusionTts(512, layer_drop=0.3, unconditioned_percentage=0.5)
+    # Test with latent aligned conditioning
+    # o = model(clip, ts, aligned_latent, cond)
+    # Test with sequence aligned conditioning
+    o = model(clip, ts, aligned_sequence, cond)

TTS/tts/layers/tortoise/random_latent_generator.py

@@ -0,0 +1,55 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def fused_leaky_relu(input, bias=None, negative_slope=0.2, scale=2**0.5):
+    if bias is not None:
+        rest_dim = [1] * (input.ndim - bias.ndim - 1)
+        return (
+            F.leaky_relu(
+                input + bias.view(1, bias.shape[0], *rest_dim),
+                negative_slope=negative_slope,
+            )
+            * scale
+        )
+    else:
+        return F.leaky_relu(input, negative_slope=0.2) * scale
+
+
+class EqualLinear(nn.Module):
+    def __init__(self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul))
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init))
+        else:
+            self.bias = None
+        self.scale = (1 / math.sqrt(in_dim)) * lr_mul
+        self.lr_mul = lr_mul
+
+    def forward(self, input):
+        out = F.linear(input, self.weight * self.scale)
+        out = fused_leaky_relu(out, self.bias * self.lr_mul)
+        return out
+
+
+class RandomLatentConverter(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.layers = nn.Sequential(
+            *[EqualLinear(channels, channels, lr_mul=0.1) for _ in range(5)], nn.Linear(channels, channels)
+        )
+        self.channels = channels
+
+    def forward(self, ref):
+        r = torch.randn(ref.shape[0], self.channels, device=ref.device)
+        y = self.layers(r)
+        return y
+
+
+if __name__ == "__main__":
+    model = RandomLatentConverter(512)
+    model(torch.randn(5, 512))

TTS/tts/layers/tortoise/tokenizer.py

@@ -0,0 +1,37 @@
+import os
+
+import torch
+from tokenizers import Tokenizer
+
+from TTS.tts.utils.text.cleaners import english_cleaners
+
+DEFAULT_VOCAB_FILE = os.path.join(
+    os.path.dirname(os.path.realpath(__file__)), "../../utils/assets/tortoise/tokenizer.json"
+)
+
+
+class VoiceBpeTokenizer:
+    def __init__(self, vocab_file=DEFAULT_VOCAB_FILE, vocab_str=None):
+        self.tokenizer = None
+        if vocab_file is not None:
+            self.tokenizer = Tokenizer.from_file(vocab_file)
+        if vocab_str is not None:
+            self.tokenizer = Tokenizer.from_str(vocab_str)
+
+    def preprocess_text(self, txt):
+        txt = english_cleaners(txt)
+        return txt
+
+    def encode(self, txt):
+        txt = self.preprocess_text(txt)
+        txt = txt.replace(" ", "[SPACE]")
+        return self.tokenizer.encode(txt).ids
+
+    def decode(self, seq):
+        if isinstance(seq, torch.Tensor):
+            seq = seq.cpu().numpy()
+        txt = self.tokenizer.decode(seq, skip_special_tokens=False).replace(" ", "")
+        txt = txt.replace("[SPACE]", " ")
+        txt = txt.replace("[STOP]", "")
+        txt = txt.replace("[UNK]", "")
+        return txt

TTS/tts/layers/tortoise/transformer.py

@@ -0,0 +1,229 @@
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import nn
+
+# helpers
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+def cast_tuple(val, depth=1):
+    if isinstance(val, list):
+        val = tuple(val)
+    return val if isinstance(val, tuple) else (val,) * depth
+
+
+def max_neg_value(t):
+    return -torch.finfo(t.dtype).max
+
+
+def stable_softmax(t, dim=-1, alpha=32**2):
+    t = t / alpha
+    t = t - torch.amax(t, dim=dim, keepdim=True).detach()
+    return (t * alpha).softmax(dim=dim)
+
+
+def route_args(router, args, depth):
+    routed_args = [(dict(), dict()) for _ in range(depth)]
+    matched_keys = [key for key in args.keys() if key in router]
+
+    for key in matched_keys:
+        val = args[key]
+        for depth, ((f_args, g_args), routes) in enumerate(zip(routed_args, router[key])):
+            new_f_args, new_g_args = map(lambda route: ({key: val} if route else {}), routes)
+            routed_args[depth] = ({**f_args, **new_f_args}, {**g_args, **new_g_args})
+    return routed_args
+
+
+# classes
+class SequentialSequence(nn.Module):
+    def __init__(self, layers, args_route={}, layer_dropout=0.0):
+        super().__init__()
+        assert all(
+            len(route) == len(layers) for route in args_route.values()
+        ), "each argument route map must have the same depth as the number of sequential layers"
+        self.layers = layers
+        self.args_route = args_route
+        self.layer_dropout = layer_dropout
+
+    def forward(self, x, **kwargs):
+        args = route_args(self.args_route, kwargs, len(self.layers))
+        layers_and_args = list(zip(self.layers, args))
+
+        for (f, g), (f_args, g_args) in layers_and_args:
+            x = x + f(x, **f_args)
+            x = x + g(x, **g_args)
+        return x
+
+
+class DivideMax(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        maxes = x.amax(dim=self.dim, keepdim=True).detach()
+        return x / maxes
+
+
+# https://arxiv.org/abs/2103.17239
+class LayerScale(nn.Module):
+    def __init__(self, dim, depth, fn):
+        super().__init__()
+        if depth <= 18:
+            init_eps = 0.1
+        elif depth > 18 and depth <= 24:
+            init_eps = 1e-5
+        else:
+            init_eps = 1e-6
+
+        scale = torch.zeros(1, 1, dim).fill_(init_eps)
+        self.scale = nn.Parameter(scale)
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        return self.fn(x, **kwargs) * self.scale
+
+
+# layer norm
+
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn, sandwich=False):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.norm_out = nn.LayerNorm(dim) if sandwich else nn.Identity()
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        x = self.norm(x)
+        x = self.fn(x, **kwargs)
+        return self.norm_out(x)
+
+
+# feed forward
+
+
+class GEGLU(nn.Module):
+    def forward(self, x):
+        x, gates = x.chunk(2, dim=-1)
+        return x * F.gelu(gates)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dropout=0.0, mult=4.0):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, dim * mult * 2),
+            GEGLU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim * mult, dim),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+# Attention
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, seq_len, causal=True, heads=8, dim_head=64, dropout=0.0):
+        super().__init__()
+        inner_dim = dim_head * heads
+        self.heads = heads
+        self.seq_len = seq_len
+        self.scale = dim_head**-0.5
+
+        self.causal = causal
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
+        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))
+
+    def forward(self, x, mask=None):
+        b, n, _, h, device = *x.shape, self.heads, x.device
+        softmax = torch.softmax
+
+        qkv = self.to_qkv(x).chunk(3, dim=-1)
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), qkv)
+
+        q = q * self.scale
+
+        dots = torch.einsum("b h i d, b h j d -> b h i j", q, k)
+        mask_value = max_neg_value(dots)
+
+        if exists(mask):
+            mask = rearrange(mask, "b j -> b () () j")
+            dots.masked_fill_(~mask, mask_value)
+            del mask
+
+        if self.causal:
+            i, j = dots.shape[-2:]
+            mask = torch.ones(i, j, device=device).triu_(j - i + 1).bool()
+            dots.masked_fill_(mask, mask_value)
+
+        attn = softmax(dots, dim=-1)
+
+        out = torch.einsum("b h i j, b h j d -> b h i d", attn, v)
+        out = rearrange(out, "b h n d -> b n (h d)")
+        out = self.to_out(out)
+        return out
+
+
+# main transformer class
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth,
+        seq_len,
+        causal=True,
+        heads=8,
+        dim_head=64,
+        ff_mult=4,
+        attn_dropout=0.0,
+        ff_dropout=0.0,
+        sparse_attn=False,
+        sandwich_norm=False,
+    ):
+        super().__init__()
+        layers = nn.ModuleList([])
+        sparse_layer = cast_tuple(sparse_attn, depth)
+
+        for ind, sparse_attn in zip(range(depth), sparse_layer):
+            attn = Attention(
+                dim,
+                causal=causal,
+                seq_len=seq_len,
+                heads=heads,
+                dim_head=dim_head,
+                dropout=attn_dropout,
+            )
+
+            ff = FeedForward(dim, mult=ff_mult, dropout=ff_dropout)
+
+            layers.append(
+                nn.ModuleList(
+                    [
+                        LayerScale(dim, ind + 1, PreNorm(dim, attn, sandwich=sandwich_norm)),
+                        LayerScale(dim, ind + 1, PreNorm(dim, ff, sandwich=sandwich_norm)),
+                    ]
+                )
+            )
+
+        execute_type = SequentialSequence
+        route_attn = ((True, False),) * depth
+        attn_route_map = {"mask": route_attn}
+
+        self.layers = execute_type(layers, args_route=attn_route_map)
+
+    def forward(self, x, **kwargs):
+        return self.layers(x, **kwargs)

TTS/tts/layers/tortoise/utils.py

@@ -0,0 +1,46 @@
+import os
+from urllib import request
+
+from tqdm import tqdm
+
+DEFAULT_MODELS_DIR = os.path.join(os.path.expanduser("~"), ".cache", "tortoise", "models")
+MODELS_DIR = os.environ.get("TORTOISE_MODELS_DIR", DEFAULT_MODELS_DIR)
+MODELS_DIR = "/data/speech_synth/models/"
+MODELS = {
+    "autoregressive.pth": "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/main/.models/autoregressive.pth",
+    "classifier.pth": "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/main/.models/classifier.pth",
+    "clvp2.pth": "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/main/.models/clvp2.pth",
+    "diffusion_decoder.pth": "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/main/.models/diffusion_decoder.pth",
+    "vocoder.pth": "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/main/.models/vocoder.pth",
+    "rlg_auto.pth": "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/main/.models/rlg_auto.pth",
+    "rlg_diffuser.pth": "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/main/.models/rlg_diffuser.pth",
+}
+
+
+def download_models(specific_models=None):
+    """
+    Call to download all the models that Tortoise uses.
+    """
+    os.makedirs(MODELS_DIR, exist_ok=True)
+    for model_name, url in MODELS.items():
+        if specific_models is not None and model_name not in specific_models:
+            continue
+        model_path = os.path.join(MODELS_DIR, model_name)
+        if os.path.exists(model_path):
+            continue
+        print(f"Downloading {model_name} from {url}...")
+        with tqdm(unit="B", unit_scale=True, unit_divisor=1024, miniters=1) as t:
+            request.urlretrieve(url, model_path, lambda nb, bs, fs, t=t: t.update(nb * bs - t.n))
+        print("Done.")
+
+
+def get_model_path(model_name, models_dir=MODELS_DIR):
+    """
+    Get path to given model, download it if it doesn't exist.
+    """
+    if model_name not in MODELS:
+        raise ValueError(f"Model {model_name} not found in available models.")
+    model_path = os.path.join(models_dir, model_name)
+    if not os.path.exists(model_path) and models_dir == MODELS_DIR:
+        download_models([model_name])
+    return model_path

TTS/tts/layers/tortoise/vocoder.py

@@ -0,0 +1,405 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Callable, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.utils.parametrize as parametrize
+
+MAX_WAV_VALUE = 32768.0
+
+
+class KernelPredictor(torch.nn.Module):
+    """Kernel predictor for the location-variable convolutions"""
+
+    def __init__(
+        self,
+        cond_channels,
+        conv_in_channels,
+        conv_out_channels,
+        conv_layers,
+        conv_kernel_size=3,
+        kpnet_hidden_channels=64,
+        kpnet_conv_size=3,
+        kpnet_dropout=0.0,
+        kpnet_nonlinear_activation="LeakyReLU",
+        kpnet_nonlinear_activation_params={"negative_slope": 0.1},
+    ):
+        """
+        Args:
+            cond_channels (int): number of channel for the conditioning sequence,
+            conv_in_channels (int): number of channel for the input sequence,
+            conv_out_channels (int): number of channel for the output sequence,
+            conv_layers (int): number of layers
+        """
+        super().__init__()
+
+        self.conv_in_channels = conv_in_channels
+        self.conv_out_channels = conv_out_channels
+        self.conv_kernel_size = conv_kernel_size
+        self.conv_layers = conv_layers
+
+        kpnet_kernel_channels = conv_in_channels * conv_out_channels * conv_kernel_size * conv_layers  # l_w
+        kpnet_bias_channels = conv_out_channels * conv_layers  # l_b
+
+        self.input_conv = nn.Sequential(
+            nn.utils.parametrizations.weight_norm(
+                nn.Conv1d(cond_channels, kpnet_hidden_channels, 5, padding=2, bias=True)
+            ),
+            getattr(nn, kpnet_nonlinear_activation)(**kpnet_nonlinear_activation_params),
+        )
+
+        self.residual_convs = nn.ModuleList()
+        padding = (kpnet_conv_size - 1) // 2
+        for _ in range(3):
+            self.residual_convs.append(
+                nn.Sequential(
+                    nn.Dropout(kpnet_dropout),
+                    nn.utils.parametrizations.weight_norm(
+                        nn.Conv1d(
+                            kpnet_hidden_channels,
+                            kpnet_hidden_channels,
+                            kpnet_conv_size,
+                            padding=padding,
+                            bias=True,
+                        )
+                    ),
+                    getattr(nn, kpnet_nonlinear_activation)(**kpnet_nonlinear_activation_params),
+                    nn.utils.parametrizations.weight_norm(
+                        nn.Conv1d(
+                            kpnet_hidden_channels,
+                            kpnet_hidden_channels,
+                            kpnet_conv_size,
+                            padding=padding,
+                            bias=True,
+                        )
+                    ),
+                    getattr(nn, kpnet_nonlinear_activation)(**kpnet_nonlinear_activation_params),
+                )
+            )
+        self.kernel_conv = nn.utils.parametrizations.weight_norm(
+            nn.Conv1d(
+                kpnet_hidden_channels,
+                kpnet_kernel_channels,
+                kpnet_conv_size,
+                padding=padding,
+                bias=True,
+            )
+        )
+        self.bias_conv = nn.utils.parametrizations.weight_norm(
+            nn.Conv1d(
+                kpnet_hidden_channels,
+                kpnet_bias_channels,
+                kpnet_conv_size,
+                padding=padding,
+                bias=True,
+            )
+        )
+
+    def forward(self, c):
+        """
+        Args:
+            c (Tensor): the conditioning sequence (batch, cond_channels, cond_length)
+        """
+        batch, _, cond_length = c.shape
+        c = self.input_conv(c)
+        for residual_conv in self.residual_convs:
+            residual_conv.to(c.device)
+            c = c + residual_conv(c)
+        k = self.kernel_conv(c)
+        b = self.bias_conv(c)
+        kernels = k.contiguous().view(
+            batch,
+            self.conv_layers,
+            self.conv_in_channels,
+            self.conv_out_channels,
+            self.conv_kernel_size,
+            cond_length,
+        )
+        bias = b.contiguous().view(
+            batch,
+            self.conv_layers,
+            self.conv_out_channels,
+            cond_length,
+        )
+
+        return kernels, bias
+
+    def remove_weight_norm(self):
+        parametrize.remove_parametrizations(self.input_conv[0], "weight")
+        parametrize.remove_parametrizations(self.kernel_conv, "weight")
+        parametrize.remove_parametrizations(self.bias_conv)
+        for block in self.residual_convs:
+            parametrize.remove_parametrizations(block[1], "weight")
+            parametrize.remove_parametrizations(block[3], "weight")
+
+
+class LVCBlock(torch.nn.Module):
+    """the location-variable convolutions"""
+
+    def __init__(
+        self,
+        in_channels,
+        cond_channels,
+        stride,
+        dilations=[1, 3, 9, 27],
+        lReLU_slope=0.2,
+        conv_kernel_size=3,
+        cond_hop_length=256,
+        kpnet_hidden_channels=64,
+        kpnet_conv_size=3,
+        kpnet_dropout=0.0,
+    ):
+        super().__init__()
+
+        self.cond_hop_length = cond_hop_length
+        self.conv_layers = len(dilations)
+        self.conv_kernel_size = conv_kernel_size
+
+        self.kernel_predictor = KernelPredictor(
+            cond_channels=cond_channels,
+            conv_in_channels=in_channels,
+            conv_out_channels=2 * in_channels,
+            conv_layers=len(dilations),
+            conv_kernel_size=conv_kernel_size,
+            kpnet_hidden_channels=kpnet_hidden_channels,
+            kpnet_conv_size=kpnet_conv_size,
+            kpnet_dropout=kpnet_dropout,
+            kpnet_nonlinear_activation_params={"negative_slope": lReLU_slope},
+        )
+
+        self.convt_pre = nn.Sequential(
+            nn.LeakyReLU(lReLU_slope),
+            nn.utils.parametrizations.weight_norm(
+                nn.ConvTranspose1d(
+                    in_channels,
+                    in_channels,
+                    2 * stride,
+                    stride=stride,
+                    padding=stride // 2 + stride % 2,
+                    output_padding=stride % 2,
+                )
+            ),
+        )
+
+        self.conv_blocks = nn.ModuleList()
+        for dilation in dilations:
+            self.conv_blocks.append(
+                nn.Sequential(
+                    nn.LeakyReLU(lReLU_slope),
+                    nn.utils.parametrizations.weight_norm(
+                        nn.Conv1d(
+                            in_channels,
+                            in_channels,
+                            conv_kernel_size,
+                            padding=dilation * (conv_kernel_size - 1) // 2,
+                            dilation=dilation,
+                        )
+                    ),
+                    nn.LeakyReLU(lReLU_slope),
+                )
+            )
+
+    def forward(self, x, c):
+        """forward propagation of the location-variable convolutions.
+        Args:
+            x (Tensor): the input sequence (batch, in_channels, in_length)
+            c (Tensor): the conditioning sequence (batch, cond_channels, cond_length)
+
+        Returns:
+            Tensor: the output sequence (batch, in_channels, in_length)
+        """
+        _, in_channels, _ = x.shape  # (B, c_g, L')
+
+        x = self.convt_pre(x)  # (B, c_g, stride * L')
+        kernels, bias = self.kernel_predictor(c)
+
+        for i, conv in enumerate(self.conv_blocks):
+            output = conv(x)  # (B, c_g, stride * L')
+
+            k = kernels[:, i, :, :, :, :]  # (B, 2 * c_g, c_g, kernel_size, cond_length)
+            b = bias[:, i, :, :]  # (B, 2 * c_g, cond_length)
+
+            output = self.location_variable_convolution(
+                output, k, b, hop_size=self.cond_hop_length
+            )  # (B, 2 * c_g, stride * L'): LVC
+            x = x + torch.sigmoid(output[:, :in_channels, :]) * torch.tanh(
+                output[:, in_channels:, :]
+            )  # (B, c_g, stride * L'): GAU
+
+        return x
+
+    def location_variable_convolution(self, x, kernel, bias, dilation=1, hop_size=256):
+        """perform location-variable convolution operation on the input sequence (x) using the local convolution kernl.
+        Time: 414 μs ± 309 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each), test on NVIDIA V100.
+        Args:
+            x (Tensor): the input sequence (batch, in_channels, in_length).
+            kernel (Tensor): the local convolution kernel (batch, in_channel, out_channels, kernel_size, kernel_length)
+            bias (Tensor): the bias for the local convolution (batch, out_channels, kernel_length)
+            dilation (int): the dilation of convolution.
+            hop_size (int): the hop_size of the conditioning sequence.
+        Returns:
+            (Tensor): the output sequence after performing local convolution. (batch, out_channels, in_length).
+        """
+        batch, _, in_length = x.shape
+        batch, _, out_channels, kernel_size, kernel_length = kernel.shape
+        assert in_length == (kernel_length * hop_size), "length of (x, kernel) is not matched"
+
+        padding = dilation * int((kernel_size - 1) / 2)
+        x = F.pad(x, (padding, padding), "constant", 0)  # (batch, in_channels, in_length + 2*padding)
+        x = x.unfold(2, hop_size + 2 * padding, hop_size)  # (batch, in_channels, kernel_length, hop_size + 2*padding)
+
+        if hop_size < dilation:
+            x = F.pad(x, (0, dilation), "constant", 0)
+        x = x.unfold(
+            3, dilation, dilation
+        )  # (batch, in_channels, kernel_length, (hop_size + 2*padding)/dilation, dilation)
+        x = x[:, :, :, :, :hop_size]
+        x = x.transpose(3, 4)  # (batch, in_channels, kernel_length, dilation, (hop_size + 2*padding)/dilation)
+        x = x.unfold(4, kernel_size, 1)  # (batch, in_channels, kernel_length, dilation, _, kernel_size)
+
+        o = torch.einsum("bildsk,biokl->bolsd", x, kernel)
+        o = o.to(memory_format=torch.channels_last_3d)
+        bias = bias.unsqueeze(-1).unsqueeze(-1).to(memory_format=torch.channels_last_3d)
+        o = o + bias
+        o = o.contiguous().view(batch, out_channels, -1)
+
+        return o
+
+    def remove_weight_norm(self):
+        self.kernel_predictor.remove_weight_norm()
+        parametrize.remove_parametrizations(self.convt_pre[1], "weight")
+        for block in self.conv_blocks:
+            parametrize.remove_parametrizations(block[1], "weight")
+
+
+class UnivNetGenerator(nn.Module):
+    """
+    UnivNet Generator
+
+    Originally from https://github.com/mindslab-ai/univnet/blob/master/model/generator.py.
+    """
+
+    def __init__(
+        self,
+        noise_dim=64,
+        channel_size=32,
+        dilations=[1, 3, 9, 27],
+        strides=[8, 8, 4],
+        lReLU_slope=0.2,
+        kpnet_conv_size=3,
+        # Below are MEL configurations options that this generator requires.
+        hop_length=256,
+        n_mel_channels=100,
+    ):
+        super(UnivNetGenerator, self).__init__()
+        self.mel_channel = n_mel_channels
+        self.noise_dim = noise_dim
+        self.hop_length = hop_length
+        channel_size = channel_size
+        kpnet_conv_size = kpnet_conv_size
+
+        self.res_stack = nn.ModuleList()
+        hop_length = 1
+        for stride in strides:
+            hop_length = stride * hop_length
+            self.res_stack.append(
+                LVCBlock(
+                    channel_size,
+                    n_mel_channels,
+                    stride=stride,
+                    dilations=dilations,
+                    lReLU_slope=lReLU_slope,
+                    cond_hop_length=hop_length,
+                    kpnet_conv_size=kpnet_conv_size,
+                )
+            )
+
+        self.conv_pre = nn.utils.parametrizations.weight_norm(
+            nn.Conv1d(noise_dim, channel_size, 7, padding=3, padding_mode="reflect")
+        )
+
+        self.conv_post = nn.Sequential(
+            nn.LeakyReLU(lReLU_slope),
+            nn.utils.parametrizations.weight_norm(nn.Conv1d(channel_size, 1, 7, padding=3, padding_mode="reflect")),
+            nn.Tanh(),
+        )
+
+    def forward(self, c, z):
+        """
+        Args:
+            c (Tensor): the conditioning sequence of mel-spectrogram (batch, mel_channels, in_length)
+            z (Tensor): the noise sequence (batch, noise_dim, in_length)
+
+        """
+        z = self.conv_pre(z)  # (B, c_g, L)
+
+        for res_block in self.res_stack:
+            res_block.to(z.device)
+            z = res_block(z, c)  # (B, c_g, L * s_0 * ... * s_i)
+
+        z = self.conv_post(z)  # (B, 1, L * 256)
+
+        return z
+
+    def eval(self, inference=False):
+        super(UnivNetGenerator, self).eval()
+        # don't remove weight norm while validation in training loop
+        if inference:
+            self.remove_weight_norm()
+
+    def remove_weight_norm(self):
+        parametrize.remove_parametrizations(self.conv_pre, "weight")
+
+        for layer in self.conv_post:
+            if len(layer.state_dict()) != 0:
+                parametrize.remove_parametrizations(layer, "weight")
+
+        for res_block in self.res_stack:
+            res_block.remove_weight_norm()
+
+    def inference(self, c, z=None):
+        # pad input mel with zeros to cut artifact
+        # see https://github.com/seungwonpark/melgan/issues/8
+        zero = torch.full((c.shape[0], self.mel_channel, 10), -11.5129).to(c.device)
+        mel = torch.cat((c, zero), dim=2)
+
+        if z is None:
+            z = torch.randn(c.shape[0], self.noise_dim, mel.size(2)).to(mel.device)
+
+        audio = self.forward(mel, z)
+        audio = audio[:, :, : -(self.hop_length * 10)]
+        audio = audio.clamp(min=-1, max=1)
+        return audio
+
+
+@dataclass
+class VocType:
+    constructor: Callable[[], nn.Module]
+    model_path: str
+    subkey: Optional[str] = None
+
+    def optionally_index(self, model_dict):
+        if self.subkey is not None:
+            return model_dict[self.subkey]
+        return model_dict
+
+
+class VocConf(Enum):
+    Univnet = VocType(UnivNetGenerator, "vocoder.pth", "model_g")
+
+
+if __name__ == "__main__":
+    model = UnivNetGenerator()
+
+    c = torch.randn(3, 100, 10)
+    z = torch.randn(3, 64, 10)
+    print(c.shape)
+
+    y = model(c, z)
+    print(y.shape)
+    assert y.shape == torch.Size([3, 1, 2560])
+
+    pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(pytorch_total_params)

TTS/tts/layers/tortoise/wav2vec_alignment.py

@@ -0,0 +1,150 @@
+import torch
+import torchaudio
+from transformers import Wav2Vec2CTCTokenizer, Wav2Vec2FeatureExtractor, Wav2Vec2ForCTC
+
+
+def max_alignment(s1, s2, skip_character="~", record=None):
+    """
+    A clever function that aligns s1 to s2 as best it can. Wherever a character from s1 is not found in s2, a '~' is
+    used to replace that character.
+
+    Finally got to use my DP skills!
+    """
+    if record is None:
+        record = {}
+    assert skip_character not in s1, f"Found the skip character {skip_character} in the provided string, {s1}"
+    if len(s1) == 0:
+        return ""
+    if len(s2) == 0:
+        return skip_character * len(s1)
+    if s1 == s2:
+        return s1
+    if s1[0] == s2[0]:
+        return s1[0] + max_alignment(s1[1:], s2[1:], skip_character, record)
+
+    take_s1_key = (len(s1), len(s2) - 1)
+    if take_s1_key in record:
+        take_s1, take_s1_score = record[take_s1_key]
+    else:
+        take_s1 = max_alignment(s1, s2[1:], skip_character, record)
+        take_s1_score = len(take_s1.replace(skip_character, ""))
+        record[take_s1_key] = (take_s1, take_s1_score)
+
+    take_s2_key = (len(s1) - 1, len(s2))
+    if take_s2_key in record:
+        take_s2, take_s2_score = record[take_s2_key]
+    else:
+        take_s2 = max_alignment(s1[1:], s2, skip_character, record)
+        take_s2_score = len(take_s2.replace(skip_character, ""))
+        record[take_s2_key] = (take_s2, take_s2_score)
+
+    return take_s1 if take_s1_score > take_s2_score else skip_character + take_s2
+
+
+class Wav2VecAlignment:
+    """
+    Uses wav2vec2 to perform audio<->text alignment.
+    """
+
+    def __init__(self, device="cuda"):
+        self.model = Wav2Vec2ForCTC.from_pretrained("jbetker/wav2vec2-large-robust-ft-libritts-voxpopuli").cpu()
+        self.feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("facebook/wav2vec2-large-960h")
+        self.tokenizer = Wav2Vec2CTCTokenizer.from_pretrained("jbetker/tacotron-symbols")
+        self.device = device
+
+    def align(self, audio, expected_text, audio_sample_rate=24000):
+        orig_len = audio.shape[-1]
+
+        with torch.no_grad():
+            self.model = self.model.to(self.device)
+            audio = audio.to(self.device)
+            audio = torchaudio.functional.resample(audio, audio_sample_rate, 16000)
+            clip_norm = (audio - audio.mean()) / torch.sqrt(audio.var() + 1e-7)
+            logits = self.model(clip_norm).logits
+            self.model = self.model.cpu()
+
+        logits = logits[0]
+        pred_string = self.tokenizer.decode(logits.argmax(-1).tolist())
+
+        fixed_expectation = max_alignment(expected_text.lower(), pred_string)
+        w2v_compression = orig_len // logits.shape[0]
+        expected_tokens = self.tokenizer.encode(fixed_expectation)
+        expected_chars = list(fixed_expectation)
+        if len(expected_tokens) == 1:
+            return [0]  # The alignment is simple; there is only one token.
+        expected_tokens.pop(0)  # The first token is a given.
+        expected_chars.pop(0)
+
+        alignments = [0]
+
+        def pop_till_you_win():
+            if len(expected_tokens) == 0:
+                return None
+            popped = expected_tokens.pop(0)
+            popped_char = expected_chars.pop(0)
+            while popped_char == "~":
+                alignments.append(-1)
+                if len(expected_tokens) == 0:
+                    return None
+                popped = expected_tokens.pop(0)
+                popped_char = expected_chars.pop(0)
+            return popped
+
+        next_expected_token = pop_till_you_win()
+        for i, logit in enumerate(logits):
+            top = logit.argmax()
+            if next_expected_token == top:
+                alignments.append(i * w2v_compression)
+                if len(expected_tokens) > 0:
+                    next_expected_token = pop_till_you_win()
+                else:
+                    break
+
+        pop_till_you_win()
+        if not (len(expected_tokens) == 0 and len(alignments) == len(expected_text)):
+            torch.save([audio, expected_text], "alignment_debug.pth")
+            assert False, (
+                "Something went wrong with the alignment algorithm. I've dumped a file, 'alignment_debug.pth' to"
+                "your current working directory. Please report this along with the file so it can get fixed."
+            )
+
+        # Now fix up alignments. Anything with -1 should be interpolated.
+        alignments.append(orig_len)  # This'll get removed but makes the algorithm below more readable.
+        for i in range(len(alignments)):
+            if alignments[i] == -1:
+                for j in range(i + 1, len(alignments)):
+                    if alignments[j] != -1:
+                        next_found_token = j
+                        break
+                for j in range(i, next_found_token):
+                    gap = alignments[next_found_token] - alignments[i - 1]
+                    alignments[j] = (j - i + 1) * gap // (next_found_token - i + 1) + alignments[i - 1]
+
+        return alignments[:-1]
+
+    def redact(self, audio, expected_text, audio_sample_rate=24000):
+        if "[" not in expected_text:
+            return audio
+        splitted = expected_text.split("[")
+        fully_split = [splitted[0]]
+        for spl in splitted[1:]:
+            assert "]" in spl, 'Every "[" character must be paired with a "]" with no nesting.'
+            fully_split.extend(spl.split("]"))
+
+        # At this point, fully_split is a list of strings, with every other string being something that should be redacted.
+        non_redacted_intervals = []
+        last_point = 0
+        for i in range(len(fully_split)):
+            if i % 2 == 0:
+                end_interval = max(0, last_point + len(fully_split[i]) - 1)
+                non_redacted_intervals.append((last_point, end_interval))
+            last_point += len(fully_split[i])
+
+        bare_text = "".join(fully_split)
+        alignments = self.align(audio, bare_text, audio_sample_rate)
+
+        output_audio = []
+        for nri in non_redacted_intervals:
+            start, stop = nri
+            output_audio.append(audio[:, alignments[start] : alignments[stop]])
+        return torch.cat(output_audio, dim=-1)

TTS/tts/layers/vits/discriminator.py

@@ -0,0 +1,89 @@
+import torch
+from torch import nn
+from torch.nn.modules.conv import Conv1d
+
+from TTS.vocoder.models.hifigan_discriminator import DiscriminatorP, MultiPeriodDiscriminator
+
+
+class DiscriminatorS(torch.nn.Module):
+    """HiFiGAN Scale Discriminator. Channel sizes are different from the original HiFiGAN.
+
+    Args:
+        use_spectral_norm (bool): if `True` swith to spectral norm instead of weight norm.
+    """
+
+    def __init__(self, use_spectral_norm=False):
+        super().__init__()
+        norm_f = nn.utils.spectral_norm if use_spectral_norm else nn.utils.parametrizations.weight_norm
+        self.convs = nn.ModuleList(
+            [
+                norm_f(Conv1d(1, 16, 15, 1, padding=7)),
+                norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)),
+                norm_f(Conv1d(64, 256, 41, 4, groups=16, padding=20)),
+                norm_f(Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
+                norm_f(Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
+                norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
+            ]
+        )
+        self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
+
+    def forward(self, x):
+        """
+        Args:
+            x (Tensor): input waveform.
+
+        Returns:
+            Tensor: discriminator scores.
+            List[Tensor]: list of features from the convolutiona layers.
+        """
+        feat = []
+        for l in self.convs:
+            x = l(x)
+            x = torch.nn.functional.leaky_relu(x, 0.1)
+            feat.append(x)
+        x = self.conv_post(x)
+        feat.append(x)
+        x = torch.flatten(x, 1, -1)
+        return x, feat
+
+
+class VitsDiscriminator(nn.Module):
+    """VITS discriminator wrapping one Scale Discriminator and a stack of Period Discriminator.
+
+    ::
+        waveform -> ScaleDiscriminator() -> scores_sd, feats_sd --> append() -> scores, feats
+               |--> MultiPeriodDiscriminator() -> scores_mpd, feats_mpd ^
+
+    Args:
+        use_spectral_norm (bool): if `True` swith to spectral norm instead of weight norm.
+    """
+
+    def __init__(self, periods=(2, 3, 5, 7, 11), use_spectral_norm=False):
+        super().__init__()
+        self.nets = nn.ModuleList()
+        self.nets.append(DiscriminatorS(use_spectral_norm=use_spectral_norm))
+        self.nets.extend([DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods])
+
+    def forward(self, x, x_hat=None):
+        """
+        Args:
+            x (Tensor): ground truth waveform.
+            x_hat (Tensor): predicted waveform.
+
+        Returns:
+            List[Tensor]: discriminator scores.
+            List[List[Tensor]]: list of list of features from each layers of each discriminator.
+        """
+        x_scores = []
+        x_hat_scores = [] if x_hat is not None else None
+        x_feats = []
+        x_hat_feats = [] if x_hat is not None else None
+        for net in self.nets:
+            x_score, x_feat = net(x)
+            x_scores.append(x_score)
+            x_feats.append(x_feat)
+            if x_hat is not None:
+                x_hat_score, x_hat_feat = net(x_hat)
+                x_hat_scores.append(x_hat_score)
+                x_hat_feats.append(x_hat_feat)
+        return x_scores, x_feats, x_hat_scores, x_hat_feats

TTS/tts/layers/vits/networks.py

@@ -0,0 +1,288 @@
+import math
+
+import torch
+from torch import nn
+
+from TTS.tts.layers.glow_tts.glow import WN
+from TTS.tts.layers.glow_tts.transformer import RelativePositionTransformer
+from TTS.tts.utils.helpers import sequence_mask
+
+LRELU_SLOPE = 0.1
+
+
+def convert_pad_shape(pad_shape):
+    l = pad_shape[::-1]
+    pad_shape = [item for sublist in l for item in sublist]
+    return pad_shape
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+class TextEncoder(nn.Module):
+    def __init__(
+        self,
+        n_vocab: int,
+        out_channels: int,
+        hidden_channels: int,
+        hidden_channels_ffn: int,
+        num_heads: int,
+        num_layers: int,
+        kernel_size: int,
+        dropout_p: float,
+        language_emb_dim: int = None,
+    ):
+        """Text Encoder for VITS model.
+
+        Args:
+            n_vocab (int): Number of characters for the embedding layer.
+            out_channels (int): Number of channels for the output.
+            hidden_channels (int): Number of channels for the hidden layers.
+            hidden_channels_ffn (int): Number of channels for the convolutional layers.
+            num_heads (int): Number of attention heads for the Transformer layers.
+            num_layers (int): Number of Transformer layers.
+            kernel_size (int): Kernel size for the FFN layers in Transformer network.
+            dropout_p (float): Dropout rate for the Transformer layers.
+        """
+        super().__init__()
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+
+        self.emb = nn.Embedding(n_vocab, hidden_channels)
+
+        nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)
+
+        if language_emb_dim:
+            hidden_channels += language_emb_dim
+
+        self.encoder = RelativePositionTransformer(
+            in_channels=hidden_channels,
+            out_channels=hidden_channels,
+            hidden_channels=hidden_channels,
+            hidden_channels_ffn=hidden_channels_ffn,
+            num_heads=num_heads,
+            num_layers=num_layers,
+            kernel_size=kernel_size,
+            dropout_p=dropout_p,
+            layer_norm_type="2",
+            rel_attn_window_size=4,
+        )
+
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(self, x, x_lengths, lang_emb=None):
+        """
+        Shapes:
+            - x: :math:`[B, T]`
+            - x_length: :math:`[B]`
+        """
+        assert x.shape[0] == x_lengths.shape[0]
+        x = self.emb(x) * math.sqrt(self.hidden_channels)  # [b, t, h]
+
+        # concat the lang emb in embedding chars
+        if lang_emb is not None:
+            x = torch.cat((x, lang_emb.transpose(2, 1).expand(x.size(0), x.size(1), -1)), dim=-1)
+
+        x = torch.transpose(x, 1, -1)  # [b, h, t]
+        x_mask = torch.unsqueeze(sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)  # [b, 1, t]
+
+        x = self.encoder(x * x_mask, x_mask)
+        stats = self.proj(x) * x_mask
+
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        return x, m, logs, x_mask
+
+
+class ResidualCouplingBlock(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        num_layers,
+        dropout_p=0,
+        cond_channels=0,
+        mean_only=False,
+    ):
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        super().__init__()
+        self.half_channels = channels // 2
+        self.mean_only = mean_only
+        # input layer
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        # coupling layers
+        self.enc = WN(
+            hidden_channels,
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            num_layers,
+            dropout_p=dropout_p,
+            c_in_channels=cond_channels,
+        )
+        # output layer
+        # Initializing last layer to 0 makes the affine coupling layers
+        # do nothing at first.  This helps with training stability
+        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+        self.post.weight.data.zero_()
+        self.post.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        """
+        Note:
+            Set `reverse` to True for inference.
+
+        Shapes:
+            - x: :math:`[B, C, T]`
+            - x_mask: :math:`[B, 1, T]`
+            - g: :math:`[B, C, 1]`
+        """
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, log_scale = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            log_scale = torch.zeros_like(m)
+
+        if not reverse:
+            x1 = m + x1 * torch.exp(log_scale) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(log_scale, [1, 2])
+            return x, logdet
+        else:
+            x1 = (x1 - m) * torch.exp(-log_scale) * x_mask
+            x = torch.cat([x0, x1], 1)
+            return x
+
+
+class ResidualCouplingBlocks(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        hidden_channels: int,
+        kernel_size: int,
+        dilation_rate: int,
+        num_layers: int,
+        num_flows=4,
+        cond_channels=0,
+    ):
+        """Redisual Coupling blocks for VITS flow layers.
+
+        Args:
+            channels (int): Number of input and output tensor channels.
+            hidden_channels (int): Number of hidden network channels.
+            kernel_size (int): Kernel size of the WaveNet layers.
+            dilation_rate (int): Dilation rate of the WaveNet layers.
+            num_layers (int): Number of the WaveNet layers.
+            num_flows (int, optional): Number of Residual Coupling blocks. Defaults to 4.
+            cond_channels (int, optional): Number of channels of the conditioning tensor. Defaults to 0.
+        """
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.num_layers = num_layers
+        self.num_flows = num_flows
+        self.cond_channels = cond_channels
+
+        self.flows = nn.ModuleList()
+        for _ in range(num_flows):
+            self.flows.append(
+                ResidualCouplingBlock(
+                    channels,
+                    hidden_channels,
+                    kernel_size,
+                    dilation_rate,
+                    num_layers,
+                    cond_channels=cond_channels,
+                    mean_only=True,
+                )
+            )
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        """
+        Note:
+            Set `reverse` to True for inference.
+
+        Shapes:
+            - x: :math:`[B, C, T]`
+            - x_mask: :math:`[B, 1, T]`
+            - g: :math:`[B, C, 1]`
+        """
+        if not reverse:
+            for flow in self.flows:
+                x, _ = flow(x, x_mask, g=g, reverse=reverse)
+                x = torch.flip(x, [1])
+        else:
+            for flow in reversed(self.flows):
+                x = torch.flip(x, [1])
+                x = flow(x, x_mask, g=g, reverse=reverse)
+        return x
+
+
+class PosteriorEncoder(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        hidden_channels: int,
+        kernel_size: int,
+        dilation_rate: int,
+        num_layers: int,
+        cond_channels=0,
+    ):
+        """Posterior Encoder of VITS model.
+
+        ::
+            x -> conv1x1() -> WaveNet() (non-causal) -> conv1x1() -> split() -> [m, s] -> sample(m, s) -> z
+
+        Args:
+            in_channels (int): Number of input tensor channels.
+            out_channels (int): Number of output tensor channels.
+            hidden_channels (int): Number of hidden channels.
+            kernel_size (int): Kernel size of the WaveNet convolution layers.
+            dilation_rate (int): Dilation rate of the WaveNet layers.
+            num_layers (int): Number of the WaveNet layers.
+            cond_channels (int, optional): Number of conditioning tensor channels. Defaults to 0.
+        """
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.num_layers = num_layers
+        self.cond_channels = cond_channels
+
+        self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
+        self.enc = WN(
+            hidden_channels, hidden_channels, kernel_size, dilation_rate, num_layers, c_in_channels=cond_channels
+        )
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(self, x, x_lengths, g=None):
+        """
+        Shapes:
+            - x: :math:`[B, C, T]`
+            - x_lengths: :math:`[B, 1]`
+            - g: :math:`[B, C, 1]`
+        """
+        x_mask = torch.unsqueeze(sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
+        x = self.pre(x) * x_mask
+        x = self.enc(x, x_mask, g=g)
+        stats = self.proj(x) * x_mask
+        mean, log_scale = torch.split(stats, self.out_channels, dim=1)
+        z = (mean + torch.randn_like(mean) * torch.exp(log_scale)) * x_mask
+        return z, mean, log_scale, x_mask

TTS/tts/layers/vits/stochastic_duration_predictor.py

@@ -0,0 +1,294 @@
+import math
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from TTS.tts.layers.generic.normalization import LayerNorm2
+from TTS.tts.layers.vits.transforms import piecewise_rational_quadratic_transform
+
+
+class DilatedDepthSeparableConv(nn.Module):
+    def __init__(self, channels, kernel_size, num_layers, dropout_p=0.0) -> torch.tensor:
+        """Dilated Depth-wise Separable Convolution module.
+
+        ::
+            x |-> DDSConv(x) -> LayerNorm(x) -> GeLU(x) -> Conv1x1(x) -> LayerNorm(x) -> GeLU(x) -> + -> o
+              |-------------------------------------------------------------------------------------^
+
+        Args:
+            channels ([type]): [description]
+            kernel_size ([type]): [description]
+            num_layers ([type]): [description]
+            dropout_p (float, optional): [description]. Defaults to 0.0.
+
+        Returns:
+            torch.tensor: Network output masked by the input sequence mask.
+        """
+        super().__init__()
+        self.num_layers = num_layers
+
+        self.convs_sep = nn.ModuleList()
+        self.convs_1x1 = nn.ModuleList()
+        self.norms_1 = nn.ModuleList()
+        self.norms_2 = nn.ModuleList()
+        for i in range(num_layers):
+            dilation = kernel_size**i
+            padding = (kernel_size * dilation - dilation) // 2
+            self.convs_sep.append(
+                nn.Conv1d(channels, channels, kernel_size, groups=channels, dilation=dilation, padding=padding)
+            )
+            self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
+            self.norms_1.append(LayerNorm2(channels))
+            self.norms_2.append(LayerNorm2(channels))
+        self.dropout = nn.Dropout(dropout_p)
+
+    def forward(self, x, x_mask, g=None):
+        """
+        Shapes:
+            - x: :math:`[B, C, T]`
+            - x_mask: :math:`[B, 1, T]`
+        """
+        if g is not None:
+            x = x + g
+        for i in range(self.num_layers):
+            y = self.convs_sep[i](x * x_mask)
+            y = self.norms_1[i](y)
+            y = F.gelu(y)
+            y = self.convs_1x1[i](y)
+            y = self.norms_2[i](y)
+            y = F.gelu(y)
+            y = self.dropout(y)
+            x = x + y
+        return x * x_mask
+
+
+class ElementwiseAffine(nn.Module):
+    """Element-wise affine transform like no-population stats BatchNorm alternative.
+
+    Args:
+        channels (int): Number of input tensor channels.
+    """
+
+    def __init__(self, channels):
+        super().__init__()
+        self.translation = nn.Parameter(torch.zeros(channels, 1))
+        self.log_scale = nn.Parameter(torch.zeros(channels, 1))
+
+    def forward(self, x, x_mask, reverse=False, **kwargs):  # pylint: disable=unused-argument
+        if not reverse:
+            y = (x * torch.exp(self.log_scale) + self.translation) * x_mask
+            logdet = torch.sum(self.log_scale * x_mask, [1, 2])
+            return y, logdet
+        x = (x - self.translation) * torch.exp(-self.log_scale) * x_mask
+        return x
+
+
+class ConvFlow(nn.Module):
+    """Dilated depth separable convolutional based spline flow.
+
+    Args:
+        in_channels (int): Number of input tensor channels.
+        hidden_channels (int): Number of in network channels.
+        kernel_size (int): Convolutional kernel size.
+        num_layers (int): Number of convolutional layers.
+        num_bins (int, optional): Number of spline bins. Defaults to 10.
+        tail_bound (float, optional): Tail bound for PRQT. Defaults to 5.0.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        kernel_size: int,
+        num_layers: int,
+        num_bins=10,
+        tail_bound=5.0,
+    ):
+        super().__init__()
+        self.num_bins = num_bins
+        self.tail_bound = tail_bound
+        self.hidden_channels = hidden_channels
+        self.half_channels = in_channels // 2
+
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.convs = DilatedDepthSeparableConv(hidden_channels, kernel_size, num_layers, dropout_p=0.0)
+        self.proj = nn.Conv1d(hidden_channels, self.half_channels * (num_bins * 3 - 1), 1)
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0)
+        h = self.convs(h, x_mask, g=g)
+        h = self.proj(h) * x_mask
+
+        b, c, t = x0.shape
+        h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2)  # [b, cx?, t] -> [b, c, t, ?]
+
+        unnormalized_widths = h[..., : self.num_bins] / math.sqrt(self.hidden_channels)
+        unnormalized_heights = h[..., self.num_bins : 2 * self.num_bins] / math.sqrt(self.hidden_channels)
+        unnormalized_derivatives = h[..., 2 * self.num_bins :]
+
+        x1, logabsdet = piecewise_rational_quadratic_transform(
+            x1,
+            unnormalized_widths,
+            unnormalized_heights,
+            unnormalized_derivatives,
+            inverse=reverse,
+            tails="linear",
+            tail_bound=self.tail_bound,
+        )
+
+        x = torch.cat([x0, x1], 1) * x_mask
+        logdet = torch.sum(logabsdet * x_mask, [1, 2])
+        if not reverse:
+            return x, logdet
+        return x
+
+
+class StochasticDurationPredictor(nn.Module):
+    """Stochastic duration predictor with Spline Flows.
+
+    It applies Variational Dequantization and Variational Data Augmentation.
+
+    Paper:
+        SDP: https://arxiv.org/pdf/2106.06103.pdf
+        Spline Flow: https://arxiv.org/abs/1906.04032
+
+    ::
+        ## Inference
+
+        x -> TextCondEncoder() -> Flow() -> dr_hat
+        noise ----------------------^
+
+        ## Training
+                                                                              |---------------------|
+        x -> TextCondEncoder() -> + -> PosteriorEncoder() -> split() -> z_u, z_v -> (d - z_u) -> concat() -> Flow() -> noise
+        d -> DurCondEncoder()  -> ^                                                    |
+        |------------------------------------------------------------------------------|
+
+    Args:
+        in_channels (int): Number of input tensor channels.
+        hidden_channels (int): Number of hidden channels.
+        kernel_size (int): Kernel size of convolutional layers.
+        dropout_p (float): Dropout rate.
+        num_flows (int, optional): Number of flow blocks. Defaults to 4.
+        cond_channels (int, optional): Number of channels of conditioning tensor. Defaults to 0.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        kernel_size: int,
+        dropout_p: float,
+        num_flows=4,
+        cond_channels=0,
+        language_emb_dim=0,
+    ):
+        super().__init__()
+
+        # add language embedding dim in the input
+        if language_emb_dim:
+            in_channels += language_emb_dim
+
+        # condition encoder text
+        self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
+        self.convs = DilatedDepthSeparableConv(hidden_channels, kernel_size, num_layers=3, dropout_p=dropout_p)
+        self.proj = nn.Conv1d(hidden_channels, hidden_channels, 1)
+
+        # posterior encoder
+        self.flows = nn.ModuleList()
+        self.flows.append(ElementwiseAffine(2))
+        self.flows += [ConvFlow(2, hidden_channels, kernel_size, num_layers=3) for _ in range(num_flows)]
+
+        # condition encoder duration
+        self.post_pre = nn.Conv1d(1, hidden_channels, 1)
+        self.post_convs = DilatedDepthSeparableConv(hidden_channels, kernel_size, num_layers=3, dropout_p=dropout_p)
+        self.post_proj = nn.Conv1d(hidden_channels, hidden_channels, 1)
+
+        # flow layers
+        self.post_flows = nn.ModuleList()
+        self.post_flows.append(ElementwiseAffine(2))
+        self.post_flows += [ConvFlow(2, hidden_channels, kernel_size, num_layers=3) for _ in range(num_flows)]
+
+        if cond_channels != 0 and cond_channels is not None:
+            self.cond = nn.Conv1d(cond_channels, hidden_channels, 1)
+
+        if language_emb_dim != 0 and language_emb_dim is not None:
+            self.cond_lang = nn.Conv1d(language_emb_dim, hidden_channels, 1)
+
+    def forward(self, x, x_mask, dr=None, g=None, lang_emb=None, reverse=False, noise_scale=1.0):
+        """
+        Shapes:
+            - x: :math:`[B, C, T]`
+            - x_mask: :math:`[B, 1, T]`
+            - dr: :math:`[B, 1, T]`
+            - g: :math:`[B, C]`
+        """
+        # condition encoder text
+        x = self.pre(x)
+        if g is not None:
+            x = x + self.cond(g)
+
+        if lang_emb is not None:
+            x = x + self.cond_lang(lang_emb)
+
+        x = self.convs(x, x_mask)
+        x = self.proj(x) * x_mask
+
+        if not reverse:
+            flows = self.flows
+            assert dr is not None
+
+            # condition encoder duration
+            h = self.post_pre(dr)
+            h = self.post_convs(h, x_mask)
+            h = self.post_proj(h) * x_mask
+            noise = torch.randn(dr.size(0), 2, dr.size(2)).to(device=x.device, dtype=x.dtype) * x_mask
+            z_q = noise
+
+            # posterior encoder
+            logdet_tot_q = 0.0
+            for idx, flow in enumerate(self.post_flows):
+                z_q, logdet_q = flow(z_q, x_mask, g=(x + h))
+                logdet_tot_q = logdet_tot_q + logdet_q
+                if idx > 0:
+                    z_q = torch.flip(z_q, [1])
+
+            z_u, z_v = torch.split(z_q, [1, 1], 1)
+            u = torch.sigmoid(z_u) * x_mask
+            z0 = (dr - u) * x_mask
+
+            # posterior encoder - neg log likelihood
+            logdet_tot_q += torch.sum((F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1, 2])
+            nll_posterior_encoder = (
+                torch.sum(-0.5 * (math.log(2 * math.pi) + (noise**2)) * x_mask, [1, 2]) - logdet_tot_q
+            )
+
+            z0 = torch.log(torch.clamp_min(z0, 1e-5)) * x_mask
+            logdet_tot = torch.sum(-z0, [1, 2])
+            z = torch.cat([z0, z_v], 1)
+
+            # flow layers
+            for idx, flow in enumerate(flows):
+                z, logdet = flow(z, x_mask, g=x, reverse=reverse)
+                logdet_tot = logdet_tot + logdet
+                if idx > 0:
+                    z = torch.flip(z, [1])
+
+            # flow layers - neg log likelihood
+            nll_flow_layers = torch.sum(0.5 * (math.log(2 * math.pi) + (z**2)) * x_mask, [1, 2]) - logdet_tot
+            return nll_flow_layers + nll_posterior_encoder
+
+        flows = list(reversed(self.flows))
+        flows = flows[:-2] + [flows[-1]]  # remove a useless vflow
+        z = torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype) * noise_scale
+        for flow in flows:
+            z = torch.flip(z, [1])
+            z = flow(z, x_mask, g=x, reverse=reverse)
+
+        z0, _ = torch.split(z, [1, 1], 1)
+        logw = z0
+        return logw

TTS/tts/layers/vits/transforms.py

@@ -0,0 +1,202 @@
+# adopted from https://github.com/bayesiains/nflows
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+DEFAULT_MIN_BIN_WIDTH = 1e-3
+DEFAULT_MIN_BIN_HEIGHT = 1e-3
+DEFAULT_MIN_DERIVATIVE = 1e-3
+
+
+def piecewise_rational_quadratic_transform(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    tails=None,
+    tail_bound=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    if tails is None:
+        spline_fn = rational_quadratic_spline
+        spline_kwargs = {}
+    else:
+        spline_fn = unconstrained_rational_quadratic_spline
+        spline_kwargs = {"tails": tails, "tail_bound": tail_bound}
+
+    outputs, logabsdet = spline_fn(
+        inputs=inputs,
+        unnormalized_widths=unnormalized_widths,
+        unnormalized_heights=unnormalized_heights,
+        unnormalized_derivatives=unnormalized_derivatives,
+        inverse=inverse,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+        **spline_kwargs,
+    )
+    return outputs, logabsdet
+
+
+def searchsorted(bin_locations, inputs, eps=1e-6):
+    bin_locations[..., -1] += eps
+    return torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1
+
+
+def unconstrained_rational_quadratic_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    tails="linear",
+    tail_bound=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
+    outside_interval_mask = ~inside_interval_mask
+
+    outputs = torch.zeros_like(inputs)
+    logabsdet = torch.zeros_like(inputs)
+
+    if tails == "linear":
+        unnormalized_derivatives = F.pad(unnormalized_derivatives, pad=(1, 1))
+        constant = np.log(np.exp(1 - min_derivative) - 1)
+        unnormalized_derivatives[..., 0] = constant
+        unnormalized_derivatives[..., -1] = constant
+
+        outputs[outside_interval_mask] = inputs[outside_interval_mask]
+        logabsdet[outside_interval_mask] = 0
+    else:
+        raise RuntimeError("{} tails are not implemented.".format(tails))
+
+    outputs[inside_interval_mask], logabsdet[inside_interval_mask] = rational_quadratic_spline(
+        inputs=inputs[inside_interval_mask],
+        unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
+        unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
+        unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],
+        inverse=inverse,
+        left=-tail_bound,
+        right=tail_bound,
+        bottom=-tail_bound,
+        top=tail_bound,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+    )
+
+    return outputs, logabsdet
+
+
+def rational_quadratic_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    left=0.0,
+    right=1.0,
+    bottom=0.0,
+    top=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    if torch.min(inputs) < left or torch.max(inputs) > right:
+        raise ValueError("Input to a transform is not within its domain")
+
+    num_bins = unnormalized_widths.shape[-1]
+
+    if min_bin_width * num_bins > 1.0:
+        raise ValueError("Minimal bin width too large for the number of bins")
+    if min_bin_height * num_bins > 1.0:
+        raise ValueError("Minimal bin height too large for the number of bins")
+
+    widths = F.softmax(unnormalized_widths, dim=-1)
+    widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
+    cumwidths = torch.cumsum(widths, dim=-1)
+    cumwidths = F.pad(cumwidths, pad=(1, 0), mode="constant", value=0.0)
+    cumwidths = (right - left) * cumwidths + left
+    cumwidths[..., 0] = left
+    cumwidths[..., -1] = right
+    widths = cumwidths[..., 1:] - cumwidths[..., :-1]
+
+    derivatives = min_derivative + F.softplus(unnormalized_derivatives)
+
+    heights = F.softmax(unnormalized_heights, dim=-1)
+    heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
+    cumheights = torch.cumsum(heights, dim=-1)
+    cumheights = F.pad(cumheights, pad=(1, 0), mode="constant", value=0.0)
+    cumheights = (top - bottom) * cumheights + bottom
+    cumheights[..., 0] = bottom
+    cumheights[..., -1] = top
+    heights = cumheights[..., 1:] - cumheights[..., :-1]
+
+    if inverse:
+        bin_idx = searchsorted(cumheights, inputs)[..., None]
+    else:
+        bin_idx = searchsorted(cumwidths, inputs)[..., None]
+
+    input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
+    input_bin_widths = widths.gather(-1, bin_idx)[..., 0]
+
+    input_cumheights = cumheights.gather(-1, bin_idx)[..., 0]
+    delta = heights / widths
+    input_delta = delta.gather(-1, bin_idx)[..., 0]
+
+    input_derivatives = derivatives.gather(-1, bin_idx)[..., 0]
+    input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0]
+
+    input_heights = heights.gather(-1, bin_idx)[..., 0]
+
+    if inverse:
+        a = (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        ) + input_heights * (input_delta - input_derivatives)
+        b = input_heights * input_derivatives - (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        )
+        c = -input_delta * (inputs - input_cumheights)
+
+        discriminant = b.pow(2) - 4 * a * c
+        assert (discriminant >= 0).all()
+
+        root = (2 * c) / (-b - torch.sqrt(discriminant))
+        outputs = root * input_bin_widths + input_cumwidths
+
+        theta_one_minus_theta = root * (1 - root)
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta) * theta_one_minus_theta
+        )
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * root.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - root).pow(2)
+        )
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, -logabsdet
+    else:
+        theta = (inputs - input_cumwidths) / input_bin_widths
+        theta_one_minus_theta = theta * (1 - theta)
+
+        numerator = input_heights * (input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta)
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta) * theta_one_minus_theta
+        )
+        outputs = input_cumheights + numerator / denominator
+
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * theta.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - theta).pow(2)
+        )
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, logabsdet

TTS/tts/layers/xtts/dvae.py

@@ -0,0 +1,393 @@
+import functools
+from math import sqrt
+
+import torch
+import torch.distributed as distributed
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+from einops import rearrange
+
+
+def default(val, d):
+    return val if val is not None else d
+
+
+def eval_decorator(fn):
+    def inner(model, *args, **kwargs):
+        was_training = model.training
+        model.eval()
+        out = fn(model, *args, **kwargs)
+        model.train(was_training)
+        return out
+
+    return inner
+
+
+def dvae_wav_to_mel(
+    wav, mel_norms_file="../experiments/clips_mel_norms.pth", mel_norms=None, device=torch.device("cpu")
+):
+    mel_stft = torchaudio.transforms.MelSpectrogram(
+        n_fft=1024,
+        hop_length=256,
+        win_length=1024,
+        power=2,
+        normalized=False,
+        sample_rate=22050,
+        f_min=0,
+        f_max=8000,
+        n_mels=80,
+        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
+
+
+class Quantize(nn.Module):
+    def __init__(self, dim, n_embed, decay=0.99, eps=1e-5, balancing_heuristic=False, new_return_order=False):
+        super().__init__()
+
+        self.dim = dim
+        self.n_embed = n_embed
+        self.decay = decay
+        self.eps = eps
+
+        self.balancing_heuristic = balancing_heuristic
+        self.codes = None
+        self.max_codes = 64000
+        self.codes_full = False
+        self.new_return_order = new_return_order
+
+        embed = torch.randn(dim, n_embed)
+        self.register_buffer("embed", embed)
+        self.register_buffer("cluster_size", torch.zeros(n_embed))
+        self.register_buffer("embed_avg", embed.clone())
+
+    def forward(self, input, return_soft_codes=False):
+        if self.balancing_heuristic and self.codes_full:
+            h = torch.histc(self.codes, bins=self.n_embed, min=0, max=self.n_embed) / len(self.codes)
+            mask = torch.logical_or(h > 0.9, h < 0.01).unsqueeze(1)
+            ep = self.embed.permute(1, 0)
+            ea = self.embed_avg.permute(1, 0)
+            rand_embed = torch.randn_like(ep) * mask
+            self.embed = (ep * ~mask + rand_embed).permute(1, 0)
+            self.embed_avg = (ea * ~mask + rand_embed).permute(1, 0)
+            self.cluster_size = self.cluster_size * ~mask.squeeze()
+            if torch.any(mask):
+                print(f"Reset {torch.sum(mask)} embedding codes.")
+                self.codes = None
+                self.codes_full = False
+
+        flatten = input.reshape(-1, self.dim)
+        dist = flatten.pow(2).sum(1, keepdim=True) - 2 * flatten @ self.embed + self.embed.pow(2).sum(0, keepdim=True)
+        soft_codes = -dist
+        _, embed_ind = soft_codes.max(1)
+        embed_onehot = F.one_hot(embed_ind, self.n_embed).type(flatten.dtype)
+        embed_ind = embed_ind.view(*input.shape[:-1])
+        quantize = self.embed_code(embed_ind)
+
+        if self.balancing_heuristic:
+            if self.codes is None:
+                self.codes = embed_ind.flatten()
+            else:
+                self.codes = torch.cat([self.codes, embed_ind.flatten()])
+                if len(self.codes) > self.max_codes:
+                    self.codes = self.codes[-self.max_codes :]
+                    self.codes_full = True
+
+        if self.training:
+            embed_onehot_sum = embed_onehot.sum(0)
+            embed_sum = flatten.transpose(0, 1) @ embed_onehot
+
+            if distributed.is_initialized() and distributed.get_world_size() > 1:
+                distributed.all_reduce(embed_onehot_sum)
+                distributed.all_reduce(embed_sum)
+
+            self.cluster_size.data.mul_(self.decay).add_(embed_onehot_sum, alpha=1 - self.decay)
+            self.embed_avg.data.mul_(self.decay).add_(embed_sum, alpha=1 - self.decay)
+            n = self.cluster_size.sum()
+            cluster_size = (self.cluster_size + self.eps) / (n + self.n_embed * self.eps) * n
+            embed_normalized = self.embed_avg / cluster_size.unsqueeze(0)
+            self.embed.data.copy_(embed_normalized)
+
+        diff = (quantize.detach() - input).pow(2).mean()
+        quantize = input + (quantize - input).detach()
+
+        if return_soft_codes:
+            return quantize, diff, embed_ind, soft_codes.view(input.shape[:-1] + (-1,))
+        elif self.new_return_order:
+            return quantize, embed_ind, diff
+        else:
+            return quantize, diff, embed_ind
+
+    def embed_code(self, embed_id):
+        return F.embedding(embed_id, self.embed.transpose(0, 1))
+
+
+# Fits a soft-discretized input to a normal-PDF across the specified dimension.
+# In other words, attempts to force the discretization function to have a mean equal utilization across all discrete
+# values with the specified expected variance.
+class DiscretizationLoss(nn.Module):
+    def __init__(self, discrete_bins, dim, expected_variance, store_past=0):
+        super().__init__()
+        self.discrete_bins = discrete_bins
+        self.dim = dim
+        self.dist = torch.distributions.Normal(0, scale=expected_variance)
+        if store_past > 0:
+            self.record_past = True
+            self.register_buffer("accumulator_index", torch.zeros(1, dtype=torch.long, device="cpu"))
+            self.register_buffer("accumulator_filled", torch.zeros(1, dtype=torch.long, device="cpu"))
+            self.register_buffer("accumulator", torch.zeros(store_past, discrete_bins))
+        else:
+            self.record_past = False
+
+    def forward(self, x):
+        other_dims = set(range(len(x.shape))) - set([self.dim])
+        averaged = x.sum(dim=tuple(other_dims)) / x.sum()
+        averaged = averaged - averaged.mean()
+
+        if self.record_past:
+            acc_count = self.accumulator.shape[0]
+            avg = averaged.detach().clone()
+            if self.accumulator_filled > 0:
+                averaged = torch.mean(self.accumulator, dim=0) * (acc_count - 1) / acc_count + averaged / acc_count
+
+            # Also push averaged into the accumulator.
+            self.accumulator[self.accumulator_index] = avg
+            self.accumulator_index += 1
+            if self.accumulator_index >= acc_count:
+                self.accumulator_index *= 0
+                if self.accumulator_filled <= 0:
+                    self.accumulator_filled += 1
+
+        return torch.sum(-self.dist.log_prob(averaged))
+
+
+class ResBlock(nn.Module):
+    def __init__(self, chan, conv, activation):
+        super().__init__()
+        self.net = nn.Sequential(
+            conv(chan, chan, 3, padding=1),
+            activation(),
+            conv(chan, chan, 3, padding=1),
+            activation(),
+            conv(chan, chan, 1),
+        )
+
+    def forward(self, x):
+        return self.net(x) + x
+
+
+class UpsampledConv(nn.Module):
+    def __init__(self, conv, *args, **kwargs):
+        super().__init__()
+        assert "stride" in kwargs.keys()
+        self.stride = kwargs["stride"]
+        del kwargs["stride"]
+        self.conv = conv(*args, **kwargs)
+
+    def forward(self, x):
+        up = nn.functional.interpolate(x, scale_factor=self.stride, mode="nearest")
+        return self.conv(up)
+
+
+# DiscreteVAE partially derived from lucidrains DALLE implementation
+# Credit: https://github.com/lucidrains/DALLE-pytorch
+class DiscreteVAE(nn.Module):
+    def __init__(
+        self,
+        positional_dims=2,
+        num_tokens=512,
+        codebook_dim=512,
+        num_layers=3,
+        num_resnet_blocks=0,
+        hidden_dim=64,
+        channels=3,
+        stride=2,
+        kernel_size=4,
+        use_transposed_convs=True,
+        encoder_norm=False,
+        activation="relu",
+        smooth_l1_loss=False,
+        straight_through=False,
+        normalization=None,  # ((0.5,) * 3, (0.5,) * 3),
+        record_codes=False,
+        discretization_loss_averaging_steps=100,
+        lr_quantizer_args={},
+    ):
+        super().__init__()
+        has_resblocks = num_resnet_blocks > 0
+
+        self.num_tokens = num_tokens
+        self.num_layers = num_layers
+        self.straight_through = straight_through
+        self.positional_dims = positional_dims
+        self.discrete_loss = DiscretizationLoss(
+            num_tokens, 2, 1 / (num_tokens * 2), discretization_loss_averaging_steps
+        )
+
+        assert positional_dims > 0 and positional_dims < 3  # This VAE only supports 1d and 2d inputs for now.
+        if positional_dims == 2:
+            conv = nn.Conv2d
+            conv_transpose = nn.ConvTranspose2d
+        else:
+            conv = nn.Conv1d
+            conv_transpose = nn.ConvTranspose1d
+        if not use_transposed_convs:
+            conv_transpose = functools.partial(UpsampledConv, conv)
+
+        if activation == "relu":
+            act = nn.ReLU
+        elif activation == "silu":
+            act = nn.SiLU
+        else:
+            assert NotImplementedError()
+
+        enc_layers = []
+        dec_layers = []
+
+        if num_layers > 0:
+            enc_chans = [hidden_dim * 2**i for i in range(num_layers)]
+            dec_chans = list(reversed(enc_chans))
+
+            enc_chans = [channels, *enc_chans]
+
+            dec_init_chan = codebook_dim if not has_resblocks else dec_chans[0]
+            dec_chans = [dec_init_chan, *dec_chans]
+
+            enc_chans_io, dec_chans_io = map(lambda t: list(zip(t[:-1], t[1:])), (enc_chans, dec_chans))
+
+            pad = (kernel_size - 1) // 2
+            for (enc_in, enc_out), (dec_in, dec_out) in zip(enc_chans_io, dec_chans_io):
+                enc_layers.append(nn.Sequential(conv(enc_in, enc_out, kernel_size, stride=stride, padding=pad), act()))
+                if encoder_norm:
+                    enc_layers.append(nn.GroupNorm(8, enc_out))
+                dec_layers.append(
+                    nn.Sequential(conv_transpose(dec_in, dec_out, kernel_size, stride=stride, padding=pad), act())
+                )
+            dec_out_chans = dec_chans[-1]
+            innermost_dim = dec_chans[0]
+        else:
+            enc_layers.append(nn.Sequential(conv(channels, hidden_dim, 1), act()))
+            dec_out_chans = hidden_dim
+            innermost_dim = hidden_dim
+
+        for _ in range(num_resnet_blocks):
+            dec_layers.insert(0, ResBlock(innermost_dim, conv, act))
+            enc_layers.append(ResBlock(innermost_dim, conv, act))
+
+        if num_resnet_blocks > 0:
+            dec_layers.insert(0, conv(codebook_dim, innermost_dim, 1))
+
+        enc_layers.append(conv(innermost_dim, codebook_dim, 1))
+        dec_layers.append(conv(dec_out_chans, channels, 1))
+
+        self.encoder = nn.Sequential(*enc_layers)
+        self.decoder = nn.Sequential(*dec_layers)
+
+        self.loss_fn = F.smooth_l1_loss if smooth_l1_loss else F.mse_loss
+        self.codebook = Quantize(codebook_dim, num_tokens, new_return_order=True)
+
+        # take care of normalization within class
+        self.normalization = normalization
+        self.record_codes = record_codes
+        if record_codes:
+            self.codes = torch.zeros((1228800,), dtype=torch.long)
+            self.code_ind = 0
+            self.total_codes = 0
+        self.internal_step = 0
+
+    def norm(self, images):
+        if not self.normalization is not None:
+            return images
+
+        means, stds = map(lambda t: torch.as_tensor(t).to(images), self.normalization)
+        arrange = "c -> () c () ()" if self.positional_dims == 2 else "c -> () c ()"
+        means, stds = map(lambda t: rearrange(t, arrange), (means, stds))
+        images = images.clone()
+        images.sub_(means).div_(stds)
+        return images
+
+    def get_debug_values(self, step, __):
+        if self.record_codes and self.total_codes > 0:
+            # Report annealing schedule
+            return {"histogram_codes": self.codes[: self.total_codes]}
+        else:
+            return {}
+
+    @torch.no_grad()
+    @eval_decorator
+    def get_codebook_indices(self, images):
+        img = self.norm(images)
+        logits = self.encoder(img).permute((0, 2, 3, 1) if len(img.shape) == 4 else (0, 2, 1))
+        sampled, codes, _ = self.codebook(logits)
+        self.log_codes(codes)
+        return codes
+
+    def decode(self, img_seq):
+        self.log_codes(img_seq)
+        if hasattr(self.codebook, "embed_code"):
+            image_embeds = self.codebook.embed_code(img_seq)
+        else:
+            image_embeds = F.embedding(img_seq, self.codebook.codebook)
+        b, n, d = image_embeds.shape
+
+        kwargs = {}
+        if self.positional_dims == 1:
+            arrange = "b n d -> b d n"
+        else:
+            h = w = int(sqrt(n))
+            arrange = "b (h w) d -> b d h w"
+            kwargs = {"h": h, "w": w}
+        image_embeds = rearrange(image_embeds, arrange, **kwargs)
+        images = [image_embeds]
+        for layer in self.decoder:
+            images.append(layer(images[-1]))
+        return images[-1], images[-2]
+
+    def infer(self, img):
+        img = self.norm(img)
+        logits = self.encoder(img).permute((0, 2, 3, 1) if len(img.shape) == 4 else (0, 2, 1))
+        sampled, codes, commitment_loss = self.codebook(logits)
+        return self.decode(codes)
+
+    # Note: This module is not meant to be run in forward() except while training. It has special logic which performs
+    # evaluation using quantized values when it detects that it is being run in eval() mode, which will be substantially
+    # more lossy (but useful for determining network performance).
+    def forward(self, img):
+        img = self.norm(img)
+        logits = self.encoder(img).permute((0, 2, 3, 1) if len(img.shape) == 4 else (0, 2, 1))
+        sampled, codes, commitment_loss = self.codebook(logits)
+        sampled = sampled.permute((0, 3, 1, 2) if len(img.shape) == 4 else (0, 2, 1))
+
+        if self.training:
+            out = sampled
+            for d in self.decoder:
+                out = d(out)
+            self.log_codes(codes)
+        else:
+            # This is non-differentiable, but gives a better idea of how the network is actually performing.
+            out, _ = self.decode(codes)
+
+        # reconstruction loss
+        recon_loss = self.loss_fn(img, out, reduction="none")
+
+        return recon_loss, commitment_loss, out
+
+    def log_codes(self, codes):
+        # This is so we can debug the distribution of codes being learned.
+        if self.record_codes and self.internal_step % 10 == 0:
+            codes = codes.flatten()
+            l = codes.shape[0]
+            i = self.code_ind if (self.codes.shape[0] - self.code_ind) > l else self.codes.shape[0] - l
+            self.codes[i : i + l] = codes.cpu()
+            self.code_ind = self.code_ind + l
+            if self.code_ind >= self.codes.shape[0]:
+                self.code_ind = 0
+            self.total_codes += 1
+        self.internal_step += 1

TTS/tts/layers/xtts/gpt.py

@@ -0,0 +1,611 @@
+# ported from: https://github.com/neonbjb/tortoise-tts
+
+import functools
+import math
+import random
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import GPT2Config
+
+from TTS.tts.layers.xtts.gpt_inference import GPT2InferenceModel
+from TTS.tts.layers.xtts.latent_encoder import ConditioningEncoder
+from TTS.tts.layers.xtts.perceiver_encoder import PerceiverResampler
+
+
+def null_position_embeddings(range, dim):
+    return torch.zeros((range.shape[0], range.shape[1], dim), device=range.device)
+
+
+class LearnedPositionEmbeddings(nn.Module):
+    def __init__(self, seq_len, model_dim, init=0.02, relative=False):
+        super().__init__()
+        # nn.Embedding
+        self.emb = torch.nn.Embedding(seq_len, model_dim)
+        # Initializing this way is standard for GPT-2
+        self.emb.weight.data.normal_(mean=0.0, std=init)
+        self.relative = relative
+        self.seq_len = seq_len
+
+    def forward(self, x):
+        sl = x.shape[1]
+        if self.relative:
+            start = random.randint(sl, self.seq_len) - sl
+            return self.emb(torch.arange(start, start + sl, device=x.device))
+        else:
+            return self.emb(torch.arange(0, sl, device=x.device))
+
+    def get_fixed_embedding(self, ind, dev):
+        return self.emb(torch.tensor([ind], device=dev)).unsqueeze(0)
+
+
+def build_hf_gpt_transformer(
+    layers,
+    model_dim,
+    heads,
+    max_mel_seq_len,
+    max_text_seq_len,
+    max_prompt_len,
+    checkpointing,
+):
+    """
+    GPT-2 implemented by the HuggingFace library.
+    """
+    from transformers import GPT2Config, GPT2Model
+
+    gpt_config = GPT2Config(
+        vocab_size=256,  # Unused.
+        n_positions=max_mel_seq_len + max_text_seq_len + max_prompt_len,
+        n_ctx=max_mel_seq_len + max_text_seq_len + max_prompt_len,
+        n_embd=model_dim,
+        n_layer=layers,
+        n_head=heads,
+        gradient_checkpointing=checkpointing,
+        use_cache=not checkpointing,
+    )
+    gpt = GPT2Model(gpt_config)
+    # Override the built in positional embeddings
+    del gpt.wpe
+    gpt.wpe = functools.partial(null_position_embeddings, dim=model_dim)
+    # Built-in token embeddings are unused.
+    del gpt.wte
+
+    mel_pos_emb = (
+        LearnedPositionEmbeddings(max_mel_seq_len, model_dim)
+        if max_mel_seq_len != -1
+        else functools.partial(null_position_embeddings, dim=model_dim)
+    )
+    text_pos_emb = (
+        LearnedPositionEmbeddings(max_text_seq_len, model_dim)
+        if max_mel_seq_len != -1
+        else functools.partial(null_position_embeddings, dim=model_dim)
+    )
+    # gpt = torch.compile(gpt, mode="reduce-overhead", fullgraph=True)
+    return gpt, mel_pos_emb, text_pos_emb, None, None
+
+
+class GPT(nn.Module):
+    def __init__(
+        self,
+        start_text_token=261,
+        stop_text_token=0,
+        layers=8,
+        model_dim=512,
+        heads=8,
+        max_text_tokens=120,
+        max_mel_tokens=250,
+        max_prompt_tokens=70,
+        max_conditioning_inputs=1,
+        code_stride_len=1024,
+        number_text_tokens=256,
+        num_audio_tokens=8194,
+        start_audio_token=8192,
+        stop_audio_token=8193,
+        train_solo_embeddings=False,
+        checkpointing=False,
+        average_conditioning_embeddings=False,
+        label_smoothing=0.0,
+        use_perceiver_resampler=False,
+        perceiver_cond_length_compression=256,
+    ):
+        """
+        Args:
+
+        """
+        super().__init__()
+
+        self.label_smoothing = label_smoothing
+        self.number_text_tokens = number_text_tokens
+        self.start_text_token = start_text_token
+        self.stop_text_token = stop_text_token
+        self.num_audio_tokens = num_audio_tokens
+        self.start_audio_token = start_audio_token
+        self.stop_audio_token = stop_audio_token
+        self.start_prompt_token = start_audio_token
+        self.stop_prompt_token = stop_audio_token
+        self.layers = layers
+        self.heads = heads
+        self.model_dim = model_dim
+        self.max_conditioning_inputs = max_conditioning_inputs
+        self.max_gen_mel_tokens = max_mel_tokens - self.max_conditioning_inputs - 2
+        self.max_mel_tokens = -1 if max_mel_tokens == -1 else max_mel_tokens + 2 + self.max_conditioning_inputs
+        self.max_text_tokens = -1 if max_text_tokens == -1 else max_text_tokens + 2
+        self.max_prompt_tokens = max_prompt_tokens
+        self.code_stride_len = code_stride_len
+        self.conditioning_encoder = ConditioningEncoder(80, model_dim, num_attn_heads=heads)
+        self.conditioning_dropout = nn.Dropout1d(0.1)
+        self.average_conditioning_embeddings = average_conditioning_embeddings
+        self.use_perceiver_resampler = use_perceiver_resampler
+        self.perceiver_cond_length_compression = perceiver_cond_length_compression
+
+        self.text_embedding = nn.Embedding(self.number_text_tokens, model_dim)
+        self.mel_embedding = nn.Embedding(self.num_audio_tokens, model_dim)
+
+        (
+            self.gpt,
+            self.mel_pos_embedding,
+            self.text_pos_embedding,
+            self.mel_layer_pos_embedding,
+            self.text_layer_pos_embedding,
+        ) = build_hf_gpt_transformer(
+            layers,
+            model_dim,
+            heads,
+            self.max_mel_tokens,
+            self.max_text_tokens,
+            self.max_prompt_tokens,
+            checkpointing,
+        )
+        if train_solo_embeddings:
+            self.mel_solo_embedding = nn.Parameter(torch.randn(1, 1, model_dim) * 0.02, requires_grad=True)
+            self.text_solo_embedding = nn.Parameter(torch.randn(1, 1, model_dim) * 0.02, requires_grad=True)
+        else:
+            self.mel_solo_embedding = 0
+            self.text_solo_embedding = 0
+
+        self.final_norm = nn.LayerNorm(model_dim)
+        self.text_head = nn.Linear(model_dim, self.number_text_tokens)
+        self.mel_head = nn.Linear(model_dim, self.num_audio_tokens)
+
+        if self.use_perceiver_resampler:
+            # XTTS v2
+            self.conditioning_perceiver = PerceiverResampler(
+                dim=model_dim,
+                depth=2,
+                dim_context=model_dim,
+                num_latents=32,
+                dim_head=64,
+                heads=8,
+                ff_mult=4,
+                use_flash_attn=False,
+            )
+        else:
+            # XTTS v1
+            self.prompt_embedding = nn.Embedding(self.num_audio_tokens, model_dim)
+            self.prompt_pos_embedding = LearnedPositionEmbeddings(24 * 9, model_dim)
+
+    def get_grad_norm_parameter_groups(self):
+        return {
+            "conditioning_encoder": list(self.conditioning_encoder.parameters()),
+            "conditioning_perceiver": list(self.conditioning_perceiver.parameters())
+            if self.use_perceiver_resampler
+            else None,
+            "gpt": list(self.gpt.parameters()),
+            "heads": list(self.text_head.parameters()) + list(self.mel_head.parameters()),
+        }
+
+    def init_gpt_for_inference(self, kv_cache=True, use_deepspeed=False):
+        seq_length = self.max_prompt_tokens + self.max_mel_tokens + self.max_text_tokens + 1
+        gpt_config = GPT2Config(
+            vocab_size=self.max_mel_tokens,
+            n_positions=seq_length,
+            n_ctx=seq_length,
+            n_embd=self.model_dim,
+            n_layer=self.layers,
+            n_head=self.heads,
+            gradient_checkpointing=False,
+            use_cache=True,
+        )
+        self.gpt_inference = GPT2InferenceModel(
+            gpt_config,
+            self.gpt,
+            self.mel_pos_embedding,
+            self.mel_embedding,
+            self.final_norm,
+            self.mel_head,
+            kv_cache=kv_cache,
+        )
+        self.gpt.wte = self.mel_embedding
+
+        if use_deepspeed:
+            import deepspeed
+
+            self.ds_engine = deepspeed.init_inference(
+                model=self.gpt_inference.half(),  # Transformers models
+                mp_size=1,  # Number of GPU
+                dtype=torch.float32,  # desired data type of output
+                replace_method="auto",  # Lets DS autmatically identify the layer to replace
+                replace_with_kernel_inject=True,  # replace the model with the kernel injector
+            )
+            self.gpt_inference = self.ds_engine.module.eval()
+
+    def set_inputs_and_targets(self, input, start_token, stop_token):
+        inp = F.pad(input, (1, 0), value=start_token)
+        tar = F.pad(input, (0, 1), value=stop_token)
+        return inp, tar
+
+    def set_mel_padding(self, mel_input_tokens, code_lengths):
+        """
+        Given mel tokens that are derived from a padded audio clip and the actual lengths of each batch element in
+        that audio clip, reformats the tokens with stop_audio_token in place of the zero padding. This is required
+        preformatting to create a working TTS model.
+        """
+        # Set padding areas within MEL (currently it is coded with the MEL code for <zero>).
+        for b in range(len(code_lengths)):
+            actual_end = code_lengths[b]
+            if actual_end < mel_input_tokens.shape[-1]:
+                mel_input_tokens[b, actual_end:] = self.stop_audio_token
+        return mel_input_tokens
+
+    def get_logits(
+        self,
+        first_inputs,
+        first_head,
+        second_inputs=None,
+        second_head=None,
+        prompt=None,
+        get_attns=False,
+        return_latent=False,
+        attn_mask_cond=None,
+        attn_mask_text=None,
+        attn_mask_mel=None,
+    ):
+        if prompt is not None:
+            offset = prompt.shape[1]
+            if second_inputs is not None:
+                emb = torch.cat([prompt, first_inputs, second_inputs], dim=1)
+            else:
+                emb = torch.cat([prompt, first_inputs], dim=1)
+
+        # with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
+        attn_mask = None
+        if attn_mask_text is not None:
+            attn_mask = torch.cat([attn_mask_text, attn_mask_mel], dim=1)
+            if prompt is not None:
+                attn_mask_cond = torch.ones(prompt.shape[0], offset, dtype=torch.bool, device=emb.device)
+                attn_mask = torch.cat([attn_mask_cond, attn_mask], dim=1)
+
+        gpt_out = self.gpt(
+            inputs_embeds=emb,
+            return_dict=True,
+            output_attentions=get_attns,
+            attention_mask=attn_mask,
+        )
+
+        if get_attns:
+            return gpt_out.attentions
+
+        enc = gpt_out.last_hidden_state[:, offset:]
+        enc = self.final_norm(enc)
+
+        if return_latent:
+            return enc[:, : first_inputs.shape[1]], enc[:, -second_inputs.shape[1] :]
+
+        first_logits = enc[:, : first_inputs.shape[1]]
+        first_logits = first_head(first_logits)
+        first_logits = first_logits.permute(0, 2, 1)
+        if second_inputs is not None:
+            second_logits = enc[:, -second_inputs.shape[1] :]
+            second_logits = second_head(second_logits)
+            second_logits = second_logits.permute(0, 2, 1)
+            return first_logits, second_logits
+        else:
+            return first_logits
+
+    def get_conditioning(self, speech_conditioning_input):
+        speech_conditioning_input = (
+            speech_conditioning_input.unsqueeze(1)
+            if len(speech_conditioning_input.shape) == 3
+            else speech_conditioning_input
+        )
+        conds = []
+        for j in range(speech_conditioning_input.shape[1]):
+            conds.append(self.conditioning_encoder(speech_conditioning_input[:, j]))
+        conds = torch.stack(conds, dim=1)
+        conds = conds.mean(dim=1)
+        return conds
+
+    def get_prompts(self, prompt_codes):
+        """
+        Create a prompt from the mel codes. This is used to condition the model on the mel codes.
+        Pad the prompt with start and stop mel tokens.
+        """
+        prompt = prompt_codes
+        if self.training:
+            lengths = []
+            # Compute the real prompt length based on the first encounter with the token 83 used for padding
+            for i in range(prompt_codes.shape[0]):
+                length = 0
+                for j in range(prompt_codes.shape[1]):
+                    if prompt_codes[i, j] == 83:
+                        break
+                    else:
+                        length += 1
+                lengths.append(length)
+
+            # prompt_len = random.randint(1, 9)  # in secs
+            prompt_len = 3
+            prompt_len = prompt_len * 24  # in frames
+            if prompt_codes.shape[-1] >= prompt_len:
+                for i in range(prompt_codes.shape[0]):
+                    if lengths[i] < prompt_len:
+                        start = 0
+                    else:
+                        start = random.randint(0, lengths[i] - prompt_len)
+                prompt = prompt_codes[:, start : start + prompt_len]
+
+        # add start and stop tokens
+        prompt = F.pad(prompt, (1, 0), value=self.start_prompt_token)
+        prompt = F.pad(prompt, (0, 1), value=self.stop_prompt_token)
+        return prompt
+
+    def get_style_emb(self, cond_input, return_latent=False):
+        """
+        cond_input: (b, 80, s) or (b, 1, 80, s)
+        conds: (b, 1024, s)
+        """
+        conds = None
+        if not return_latent:
+            if cond_input.ndim == 4:
+                cond_input = cond_input.squeeze(1)
+            conds = self.conditioning_encoder(cond_input)  # (b, d, s)
+            if self.use_perceiver_resampler:
+                conds = self.conditioning_perceiver(conds.permute(0, 2, 1)).transpose(1, 2)  # (b, d, 32)
+        else:
+            # already computed
+            conds = cond_input.unsqueeze(1)
+        return conds
+
+    def forward(
+        self,
+        text_inputs,
+        text_lengths,
+        audio_codes,
+        wav_lengths,
+        cond_mels=None,
+        cond_idxs=None,
+        cond_lens=None,
+        cond_latents=None,
+        return_attentions=False,
+        return_latent=False,
+    ):
+        """
+        Forward pass that uses both text and voice in either text conditioning mode or voice conditioning mode
+        (actuated by `text_first`).
+
+        text_inputs: long tensor, (b,t)
+        text_lengths: long tensor, (b,)
+        mel_inputs:  long tensor, (b,m)
+        wav_lengths: long tensor, (b,)
+        cond_mels: MEL float tensor, (b, 1, 80,s)
+        cond_idxs: cond start and end indexs, (b, 2)
+
+        If return_attentions is specified, only logits are returned.
+        If return_latent is specified, loss & logits are not computed or returned. Only the predicted latents are returned.
+        """
+        # ❗ FIXIT
+        if self.max_conditioning_inputs == 0:
+            assert cond_mels is None, " ❗ cond_mels is not None, but max_conditioning_inputs == 0"
+
+        max_text_len = text_lengths.max()
+        code_lengths = torch.ceil(wav_lengths / self.code_stride_len).long() + 3
+
+        if cond_lens is not None:
+            if self.use_perceiver_resampler:
+                cond_lens = cond_lens // self.perceiver_cond_length_compression
+            else:
+                cond_lens = cond_lens // self.code_stride_len
+
+        if cond_idxs is not None:
+            # recompute cond idxs for mel lengths
+            for idx in range(cond_idxs.size(0)):
+                if self.use_perceiver_resampler:
+                    cond_idxs[idx] = cond_idxs[idx] // self.perceiver_cond_length_compression
+                else:
+                    cond_idxs[idx] = cond_idxs[idx] // self.code_stride_len
+
+        # ensure that the cond_mel does not have padding
+        # if cond_lens is not None and cond_idxs is None:
+        #     min_cond_len = torch.min(cond_lens)
+        #     cond_mels = cond_mels[:, :, :, :min_cond_len]
+
+        # If len(codes) + 3 is larger than maxiumum allowed length, we truncate the codes.
+        max_mel_len = code_lengths.max()
+
+        if max_mel_len > audio_codes.shape[-1]:
+            audio_codes = F.pad(audio_codes, (0, max_mel_len - audio_codes.shape[-1]))
+
+        # 💖 Lovely assertions
+        assert (
+            max_mel_len <= audio_codes.shape[-1]
+        ), f" ❗ max_mel_len ({max_mel_len}) > audio_codes.shape[-1] ({audio_codes.shape[-1]})"
+        assert (
+            max_text_len <= text_inputs.shape[-1]
+        ), f" ❗ max_text_len ({max_text_len}) > text_inputs.shape[-1] ({text_inputs.shape[-1]})"
+
+        # Append stop token to text inputs
+        text_inputs = F.pad(text_inputs[:, :max_text_len], (0, 1), value=self.stop_text_token)
+
+        # Append silence token to mel codes
+        audio_codes = F.pad(audio_codes[:, :max_mel_len], (0, 1), value=self.stop_audio_token)
+
+        # Pad mel codes with stop_audio_token
+        audio_codes = self.set_mel_padding(
+            audio_codes, code_lengths - 3
+        )  # -3 to get the real code lengths without consider start and stop tokens that was not added yet
+
+        # Build input and target tensors
+        # Prepend start token to inputs and append stop token to targets
+        text_inputs, text_targets = self.set_inputs_and_targets(
+            text_inputs, self.start_text_token, self.stop_text_token
+        )
+        audio_codes, mel_targets = self.set_inputs_and_targets(
+            audio_codes, self.start_audio_token, self.stop_audio_token
+        )
+
+        # Set attn_mask
+        attn_mask_cond = None
+        attn_mask_text = None
+        attn_mask_mel = None
+        if not return_latent:
+            attn_mask_cond = torch.ones(
+                cond_mels.shape[0],
+                cond_mels.shape[-1],
+                dtype=torch.bool,
+                device=text_inputs.device,
+            )
+            attn_mask_text = torch.ones(
+                text_inputs.shape[0],
+                text_inputs.shape[1],
+                dtype=torch.bool,
+                device=text_inputs.device,
+            )
+            attn_mask_mel = torch.ones(
+                audio_codes.shape[0],
+                audio_codes.shape[1],
+                dtype=torch.bool,
+                device=audio_codes.device,
+            )
+
+            if cond_idxs is not None:
+                # use masking approach
+                for idx, r in enumerate(cond_idxs):
+                    l = r[1] - r[0]
+                    attn_mask_cond[idx, l:] = 0.0
+            elif cond_lens is not None:
+                for idx, l in enumerate(cond_lens):
+                    attn_mask_cond[idx, l:] = 0.0
+
+            for idx, l in enumerate(text_lengths):
+                attn_mask_text[idx, l + 1 :] = 0.0
+
+            for idx, l in enumerate(code_lengths):
+                attn_mask_mel[idx, l + 1 :] = 0.0
+
+        # Compute text embeddings + positional embeddings
+        text_emb = self.text_embedding(text_inputs) + self.text_pos_embedding(text_inputs)
+
+        # Compute mel embeddings + positional embeddings
+        mel_emb = self.mel_embedding(audio_codes) + self.mel_pos_embedding(audio_codes)
+
+        # Compute speech conditioning input
+        if cond_latents is None:
+            cond_latents = self.get_style_emb(cond_mels).transpose(1, 2)
+
+        # Get logits
+        sub = -5  # don't ask me why 😄
+        if self.training:
+            sub = -1
+
+        text_logits, mel_logits = self.get_logits(
+            text_emb,
+            self.text_head,
+            mel_emb,
+            self.mel_head,
+            prompt=cond_latents,
+            get_attns=return_attentions,
+            return_latent=return_latent,
+            attn_mask_cond=attn_mask_cond,
+            attn_mask_text=attn_mask_text,
+            attn_mask_mel=attn_mask_mel,
+        )
+        if return_latent:
+            return mel_logits[:, :sub]  # sub to prevent bla.
+
+        if return_attentions:
+            return mel_logits
+
+        # Set paddings to -1 to ignore them in loss
+        for idx, l in enumerate(text_lengths):
+            text_targets[idx, l + 1 :] = -1
+
+        for idx, l in enumerate(code_lengths):
+            mel_targets[idx, l + 1 :] = -1
+
+        # check if stoptoken is in every row of mel_targets
+        assert (mel_targets == self.stop_audio_token).sum() >= mel_targets.shape[
+            0
+        ], f" ❗ mel_targets does not contain stop token ({self.stop_audio_token}) in every row."
+
+        # ignore the loss for the segment used for conditioning
+        # coin flip for the segment to be ignored
+        if cond_idxs is not None:
+            cond_start = cond_idxs[idx, 0]
+            cond_end = cond_idxs[idx, 1]
+            mel_targets[idx, cond_start:cond_end] = -1
+
+        # Compute losses
+        loss_text = F.cross_entropy(
+            text_logits, text_targets.long(), ignore_index=-1, label_smoothing=self.label_smoothing
+        )
+        loss_mel = F.cross_entropy(
+            mel_logits, mel_targets.long(), ignore_index=-1, label_smoothing=self.label_smoothing
+        )
+        return loss_text.mean(), loss_mel.mean(), mel_logits
+
+    def inference(self, cond_latents, text_inputs, **hf_generate_kwargs):
+        self.compute_embeddings(cond_latents, text_inputs)
+        return self.generate(cond_latents, text_inputs, **hf_generate_kwargs)
+
+    def compute_embeddings(
+        self,
+        cond_latents,
+        text_inputs,
+    ):
+        text_inputs = F.pad(text_inputs, (0, 1), value=self.stop_text_token)
+        text_inputs = F.pad(text_inputs, (1, 0), value=self.start_text_token)
+        emb = self.text_embedding(text_inputs) + self.text_pos_embedding(text_inputs)
+        emb = torch.cat([cond_latents, emb], dim=1)
+        self.gpt_inference.store_prefix_emb(emb)
+        gpt_inputs = torch.full(
+            (
+                emb.shape[0],
+                emb.shape[1] + 1,  # +1 for the start_audio_token
+            ),
+            fill_value=1,
+            dtype=torch.long,
+            device=text_inputs.device,
+        )
+        gpt_inputs[:, -1] = self.start_audio_token
+        return gpt_inputs
+
+    def generate(
+        self,
+        cond_latents,
+        text_inputs,
+        **hf_generate_kwargs,
+    ):
+        gpt_inputs = self.compute_embeddings(cond_latents, text_inputs)
+        gen = self.gpt_inference.generate(
+            gpt_inputs,
+            bos_token_id=self.start_audio_token,
+            pad_token_id=self.stop_audio_token,
+            eos_token_id=self.stop_audio_token,
+            max_length=self.max_gen_mel_tokens + gpt_inputs.shape[-1],
+            **hf_generate_kwargs,
+        )
+        if "return_dict_in_generate" in hf_generate_kwargs:
+            return gen.sequences[:, gpt_inputs.shape[1] :], gen
+        return gen[:, gpt_inputs.shape[1] :]
+
+    def get_generator(self, fake_inputs, **hf_generate_kwargs):
+        return self.gpt_inference.generate_stream(
+            fake_inputs,
+            bos_token_id=self.start_audio_token,
+            pad_token_id=self.stop_audio_token,
+            eos_token_id=self.stop_audio_token,
+            max_length=self.max_gen_mel_tokens + fake_inputs.shape[-1],
+            do_stream=True,
+            **hf_generate_kwargs,
+        )

TTS/tts/layers/xtts/gpt_inference.py

@@ -0,0 +1,136 @@
+import math
+
+import torch
+from torch import nn
+from transformers import GPT2PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+
+
+class GPT2InferenceModel(GPT2PreTrainedModel):
+    """Override GPT2LMHeadModel to allow for prefix conditioning."""
+
+    def __init__(self, config, gpt, pos_emb, embeddings, norm, linear, kv_cache):
+        super().__init__(config)
+        self.transformer = gpt
+        self.pos_embedding = pos_emb
+        self.embeddings = embeddings
+        self.final_norm = norm
+        self.lm_head = nn.Sequential(norm, linear)
+        self.kv_cache = kv_cache
+
+    def store_prefix_emb(self, prefix_emb):
+        self.cached_prefix_emb = prefix_emb
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        token_type_ids = kwargs.get("token_type_ids", None)  # usually None
+        if not self.kv_cache:
+            past_key_values = None
+
+        # only last token for inputs_ids if past is defined in kwargs
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
+
+        attention_mask = kwargs.get("attention_mask", None)
+        position_ids = kwargs.get("position_ids", None)
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values is not None:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+        else:
+            position_ids = None
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "use_cache": kwargs.get("use_cache"),
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
+            "token_type_ids": token_type_ids,
+        }
+
+    def forward(
+        self,
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        assert self.cached_prefix_emb is not None
+        assert inputs_embeds is None  # Not supported by this inference model.
+        assert labels is None  # Training not supported by this inference model.
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # assert len(past_key_values) + len(input_ids) == attention_mask.shape[1]
+
+        # Create embedding
+        prefix_len = self.cached_prefix_emb.shape[1]
+        if input_ids.shape[1] != 1:
+            gen_inputs = input_ids[:, prefix_len:]
+            gen_emb = self.embeddings(gen_inputs)
+            gen_emb = gen_emb + self.pos_embedding(gen_emb)
+            if self.cached_prefix_emb.shape[0] != gen_emb.shape[0]:
+                prefix_emb = self.cached_prefix_emb.repeat_interleave(
+                    gen_emb.shape[0] // self.cached_prefix_emb.shape[0], 0
+                )
+            else:
+                prefix_emb = self.cached_prefix_emb.to(gen_emb.dtype)
+            emb = torch.cat([prefix_emb, gen_emb], dim=1)
+        else:
+            emb = self.embeddings(input_ids)
+            emb = emb + self.pos_embedding.get_fixed_embedding(
+                attention_mask.shape[1] - (prefix_len + 1), attention_mask.device
+            )
+        transformer_outputs = self.transformer(
+            inputs_embeds=emb,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+        lm_logits = self.lm_head(hidden_states)
+
+        if not return_dict:
+            return (lm_logits,) + transformer_outputs[1:]
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=None,
+            logits=lm_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+            cross_attentions=transformer_outputs.cross_attentions,
+        )
+
+    @staticmethod
+    def _reorder_cache(past, beam_idx):
+        """
+        This function is used to re-order the :obj:`past_key_values` cache if
+        :meth:`~transformers.PreTrainedModel.beam_search` or :meth:`~transformers.PreTrainedModel.beam_sample` is
+        called. This is required to match :obj:`past_key_values` with the correct beam_idx at every generation step.
+        """
+        return tuple(
+            tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
+            for layer_past in past
+        )

TTS/tts/layers/xtts/hifigan_decoder.py

@@ -0,0 +1,732 @@
+import torch
+import torchaudio
+from torch import nn
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn import functional as F
+from torch.nn.utils.parametrizations import weight_norm
+from torch.nn.utils.parametrize import remove_parametrizations
+
+from TTS.utils.io import load_fsspec
+
+LRELU_SLOPE = 0.1
+
+
+def get_padding(k, d):
+    return int((k * d - d) / 2)
+
+
+class ResBlock1(torch.nn.Module):
+    """Residual Block Type 1. It has 3 convolutional layers in each convolutional block.
+
+    Network::
+
+        x -> lrelu -> conv1_1 -> conv1_2 -> conv1_3 -> z -> lrelu -> conv2_1 -> conv2_2 -> conv2_3 -> o -> + -> o
+        |--------------------------------------------------------------------------------------------------|
+
+
+    Args:
+        channels (int): number of hidden channels for the convolutional layers.
+        kernel_size (int): size of the convolution filter in each layer.
+        dilations (list): list of dilation value for each conv layer in a block.
+    """
+
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super().__init__()
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[2],
+                        padding=get_padding(kernel_size, dilation[2]),
+                    )
+                ),
+            ]
+        )
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+            ]
+        )
+
+    def forward(self, x):
+        """
+        Args:
+            x (Tensor): input tensor.
+        Returns:
+            Tensor: output tensor.
+        Shapes:
+            x: [B, C, T]
+        """
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_parametrizations(l, "weight")
+        for l in self.convs2:
+            remove_parametrizations(l, "weight")
+
+
+class ResBlock2(torch.nn.Module):
+    """Residual Block Type 2. It has 1 convolutional layers in each convolutional block.
+
+    Network::
+
+        x -> lrelu -> conv1-> -> z -> lrelu -> conv2-> o -> + -> o
+        |---------------------------------------------------|
+
+
+    Args:
+        channels (int): number of hidden channels for the convolutional layers.
+        kernel_size (int): size of the convolution filter in each layer.
+        dilations (list): list of dilation value for each conv layer in a block.
+    """
+
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
+        super().__init__()
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+            ]
+        )
+
+    def forward(self, x):
+        for c in self.convs:
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            xt = c(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_parametrizations(l, "weight")
+
+
+class HifiganGenerator(torch.nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        resblock_type,
+        resblock_dilation_sizes,
+        resblock_kernel_sizes,
+        upsample_kernel_sizes,
+        upsample_initial_channel,
+        upsample_factors,
+        inference_padding=5,
+        cond_channels=0,
+        conv_pre_weight_norm=True,
+        conv_post_weight_norm=True,
+        conv_post_bias=True,
+        cond_in_each_up_layer=False,
+    ):
+        r"""HiFiGAN Generator with Multi-Receptive Field Fusion (MRF)
+
+        Network:
+            x -> lrelu -> upsampling_layer -> resblock1_k1x1 -> z1 -> + -> z_sum / #resblocks -> lrelu -> conv_post_7x1 -> tanh -> o
+                                                 ..          -> zI ---|
+                                              resblockN_kNx1 -> zN ---'
+
+        Args:
+            in_channels (int): number of input tensor channels.
+            out_channels (int): number of output tensor channels.
+            resblock_type (str): type of the `ResBlock`. '1' or '2'.
+            resblock_dilation_sizes (List[List[int]]): list of dilation values in each layer of a `ResBlock`.
+            resblock_kernel_sizes (List[int]): list of kernel sizes for each `ResBlock`.
+            upsample_kernel_sizes (List[int]): list of kernel sizes for each transposed convolution.
+            upsample_initial_channel (int): number of channels for the first upsampling layer. This is divided by 2
+                for each consecutive upsampling layer.
+            upsample_factors (List[int]): upsampling factors (stride) for each upsampling layer.
+            inference_padding (int): constant padding applied to the input at inference time. Defaults to 5.
+        """
+        super().__init__()
+        self.inference_padding = inference_padding
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_factors)
+        self.cond_in_each_up_layer = cond_in_each_up_layer
+
+        # initial upsampling layers
+        self.conv_pre = weight_norm(Conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3))
+        resblock = ResBlock1 if resblock_type == "1" else ResBlock2
+        # upsampling layers
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_factors, upsample_kernel_sizes)):
+            self.ups.append(
+                weight_norm(
+                    ConvTranspose1d(
+                        upsample_initial_channel // (2**i),
+                        upsample_initial_channel // (2 ** (i + 1)),
+                        k,
+                        u,
+                        padding=(k - u) // 2,
+                    )
+                )
+            )
+        # MRF blocks
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel // (2 ** (i + 1))
+            for _, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
+                self.resblocks.append(resblock(ch, k, d))
+        # post convolution layer
+        self.conv_post = weight_norm(Conv1d(ch, out_channels, 7, 1, padding=3, bias=conv_post_bias))
+        if cond_channels > 0:
+            self.cond_layer = nn.Conv1d(cond_channels, upsample_initial_channel, 1)
+
+        if not conv_pre_weight_norm:
+            remove_parametrizations(self.conv_pre, "weight")
+
+        if not conv_post_weight_norm:
+            remove_parametrizations(self.conv_post, "weight")
+
+        if self.cond_in_each_up_layer:
+            self.conds = nn.ModuleList()
+            for i in range(len(self.ups)):
+                ch = upsample_initial_channel // (2 ** (i + 1))
+                self.conds.append(nn.Conv1d(cond_channels, ch, 1))
+
+    def forward(self, x, g=None):
+        """
+        Args:
+            x (Tensor): feature input tensor.
+            g (Tensor): global conditioning input tensor.
+
+        Returns:
+            Tensor: output waveform.
+
+        Shapes:
+            x: [B, C, T]
+            Tensor: [B, 1, T]
+        """
+        o = self.conv_pre(x)
+        if hasattr(self, "cond_layer"):
+            o = o + self.cond_layer(g)
+        for i in range(self.num_upsamples):
+            o = F.leaky_relu(o, LRELU_SLOPE)
+            o = self.ups[i](o)
+
+            if self.cond_in_each_up_layer:
+                o = o + self.conds[i](g)
+
+            z_sum = None
+            for j in range(self.num_kernels):
+                if z_sum is None:
+                    z_sum = self.resblocks[i * self.num_kernels + j](o)
+                else:
+                    z_sum += self.resblocks[i * self.num_kernels + j](o)
+            o = z_sum / self.num_kernels
+        o = F.leaky_relu(o)
+        o = self.conv_post(o)
+        o = torch.tanh(o)
+        return o
+
+    @torch.no_grad()
+    def inference(self, c):
+        """
+        Args:
+            x (Tensor): conditioning input tensor.
+
+        Returns:
+            Tensor: output waveform.
+
+        Shapes:
+            x: [B, C, T]
+            Tensor: [B, 1, T]
+        """
+        c = c.to(self.conv_pre.weight.device)
+        c = torch.nn.functional.pad(c, (self.inference_padding, self.inference_padding), "replicate")
+        return self.forward(c)
+
+    def remove_weight_norm(self):
+        print("Removing weight norm...")
+        for l in self.ups:
+            remove_parametrizations(l, "weight")
+        for l in self.resblocks:
+            l.remove_weight_norm()
+        remove_parametrizations(self.conv_pre, "weight")
+        remove_parametrizations(self.conv_post, "weight")
+
+    def load_checkpoint(
+        self, config, checkpoint_path, eval=False, cache=False
+    ):  # pylint: disable=unused-argument, redefined-builtin
+        state = torch.load(checkpoint_path, map_location=torch.device("cpu"))
+        self.load_state_dict(state["model"])
+        if eval:
+            self.eval()
+            assert not self.training
+            self.remove_weight_norm()
+
+
+class SELayer(nn.Module):
+    def __init__(self, channel, reduction=8):
+        super(SELayer, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(channel, channel // reduction),
+            nn.ReLU(inplace=True),
+            nn.Linear(channel // reduction, channel),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        b, c, _, _ = x.size()
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1, 1)
+        return x * y
+
+
+class SEBasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, reduction=8):
+        super(SEBasicBlock, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.se = SELayer(planes, reduction)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.relu(out)
+        out = self.bn1(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.se(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+        return out
+
+
+def set_init_dict(model_dict, checkpoint_state, c):
+    # Partial initialization: if there is a mismatch with new and old layer, it is skipped.
+    for k, v in checkpoint_state.items():
+        if k not in model_dict:
+            print(" | > Layer missing in the model definition: {}".format(k))
+    # 1. filter out unnecessary keys
+    pretrained_dict = {k: v for k, v in checkpoint_state.items() if k in model_dict}
+    # 2. filter out different size layers
+    pretrained_dict = {k: v for k, v in pretrained_dict.items() if v.numel() == model_dict[k].numel()}
+    # 3. skip reinit layers
+    if c.has("reinit_layers") and c.reinit_layers is not None:
+        for reinit_layer_name in c.reinit_layers:
+            pretrained_dict = {k: v for k, v in pretrained_dict.items() if reinit_layer_name not in k}
+    # 4. overwrite entries in the existing state dict
+    model_dict.update(pretrained_dict)
+    print(" | > {} / {} layers are restored.".format(len(pretrained_dict), len(model_dict)))
+    return model_dict
+
+
+class PreEmphasis(nn.Module):
+    def __init__(self, coefficient=0.97):
+        super().__init__()
+        self.coefficient = coefficient
+        self.register_buffer("filter", torch.FloatTensor([-self.coefficient, 1.0]).unsqueeze(0).unsqueeze(0))
+
+    def forward(self, x):
+        assert len(x.size()) == 2
+
+        x = torch.nn.functional.pad(x.unsqueeze(1), (1, 0), "reflect")
+        return torch.nn.functional.conv1d(x, self.filter).squeeze(1)
+
+
+class ResNetSpeakerEncoder(nn.Module):
+    """This is copied from 🐸TTS to remove it from the dependencies."""
+
+    # pylint: disable=W0102
+    def __init__(
+        self,
+        input_dim=64,
+        proj_dim=512,
+        layers=[3, 4, 6, 3],
+        num_filters=[32, 64, 128, 256],
+        encoder_type="ASP",
+        log_input=False,
+        use_torch_spec=False,
+        audio_config=None,
+    ):
+        super(ResNetSpeakerEncoder, self).__init__()
+
+        self.encoder_type = encoder_type
+        self.input_dim = input_dim
+        self.log_input = log_input
+        self.use_torch_spec = use_torch_spec
+        self.audio_config = audio_config
+        self.proj_dim = proj_dim
+
+        self.conv1 = nn.Conv2d(1, num_filters[0], kernel_size=3, stride=1, padding=1)
+        self.relu = nn.ReLU(inplace=True)
+        self.bn1 = nn.BatchNorm2d(num_filters[0])
+
+        self.inplanes = num_filters[0]
+        self.layer1 = self.create_layer(SEBasicBlock, num_filters[0], layers[0])
+        self.layer2 = self.create_layer(SEBasicBlock, num_filters[1], layers[1], stride=(2, 2))
+        self.layer3 = self.create_layer(SEBasicBlock, num_filters[2], layers[2], stride=(2, 2))
+        self.layer4 = self.create_layer(SEBasicBlock, num_filters[3], layers[3], stride=(2, 2))
+
+        self.instancenorm = nn.InstanceNorm1d(input_dim)
+
+        if self.use_torch_spec:
+            self.torch_spec = torch.nn.Sequential(
+                PreEmphasis(audio_config["preemphasis"]),
+                torchaudio.transforms.MelSpectrogram(
+                    sample_rate=audio_config["sample_rate"],
+                    n_fft=audio_config["fft_size"],
+                    win_length=audio_config["win_length"],
+                    hop_length=audio_config["hop_length"],
+                    window_fn=torch.hamming_window,
+                    n_mels=audio_config["num_mels"],
+                ),
+            )
+
+        else:
+            self.torch_spec = None
+
+        outmap_size = int(self.input_dim / 8)
+
+        self.attention = nn.Sequential(
+            nn.Conv1d(num_filters[3] * outmap_size, 128, kernel_size=1),
+            nn.ReLU(),
+            nn.BatchNorm1d(128),
+            nn.Conv1d(128, num_filters[3] * outmap_size, kernel_size=1),
+            nn.Softmax(dim=2),
+        )
+
+        if self.encoder_type == "SAP":
+            out_dim = num_filters[3] * outmap_size
+        elif self.encoder_type == "ASP":
+            out_dim = num_filters[3] * outmap_size * 2
+        else:
+            raise ValueError("Undefined encoder")
+
+        self.fc = nn.Linear(out_dim, proj_dim)
+
+        self._init_layers()
+
+    def _init_layers(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def create_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    # pylint: disable=R0201
+    def new_parameter(self, *size):
+        out = nn.Parameter(torch.FloatTensor(*size))
+        nn.init.xavier_normal_(out)
+        return out
+
+    def forward(self, x, l2_norm=False):
+        """Forward pass of the model.
+
+        Args:
+            x (Tensor): Raw waveform signal or spectrogram frames. If input is a waveform, `torch_spec` must be `True`
+                to compute the spectrogram on-the-fly.
+            l2_norm (bool): Whether to L2-normalize the outputs.
+
+        Shapes:
+            - x: :math:`(N, 1, T_{in})` or :math:`(N, D_{spec}, T_{in})`
+        """
+        x.squeeze_(1)
+        # if you torch spec compute it otherwise use the mel spec computed by the AP
+        if self.use_torch_spec:
+            x = self.torch_spec(x)
+
+        if self.log_input:
+            x = (x + 1e-6).log()
+        x = self.instancenorm(x).unsqueeze(1)
+
+        x = self.conv1(x)
+        x = self.relu(x)
+        x = self.bn1(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = x.reshape(x.size()[0], -1, x.size()[-1])
+
+        w = self.attention(x)
+
+        if self.encoder_type == "SAP":
+            x = torch.sum(x * w, dim=2)
+        elif self.encoder_type == "ASP":
+            mu = torch.sum(x * w, dim=2)
+            sg = torch.sqrt((torch.sum((x**2) * w, dim=2) - mu**2).clamp(min=1e-5))
+            x = torch.cat((mu, sg), 1)
+
+        x = x.view(x.size()[0], -1)
+        x = self.fc(x)
+
+        if l2_norm:
+            x = torch.nn.functional.normalize(x, p=2, dim=1)
+        return x
+
+    def load_checkpoint(
+        self,
+        checkpoint_path: str,
+        eval: bool = False,
+        use_cuda: bool = False,
+        criterion=None,
+        cache=False,
+    ):
+        state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"), cache=cache)
+        try:
+            self.load_state_dict(state["model"])
+            print(" > Model fully restored. ")
+        except (KeyError, RuntimeError) as error:
+            # If eval raise the error
+            if eval:
+                raise error
+
+            print(" > Partial model initialization.")
+            model_dict = self.state_dict()
+            model_dict = set_init_dict(model_dict, state["model"])
+            self.load_state_dict(model_dict)
+            del model_dict
+
+        # load the criterion for restore_path
+        if criterion is not None and "criterion" in state:
+            try:
+                criterion.load_state_dict(state["criterion"])
+            except (KeyError, RuntimeError) as error:
+                print(" > Criterion load ignored because of:", error)
+
+        if use_cuda:
+            self.cuda()
+            if criterion is not None:
+                criterion = criterion.cuda()
+
+        if eval:
+            self.eval()
+            assert not self.training
+
+        if not eval:
+            return criterion, state["step"]
+        return criterion
+
+
+class HifiDecoder(torch.nn.Module):
+    def __init__(
+        self,
+        input_sample_rate=22050,
+        output_sample_rate=24000,
+        output_hop_length=256,
+        ar_mel_length_compression=1024,
+        decoder_input_dim=1024,
+        resblock_type_decoder="1",
+        resblock_dilation_sizes_decoder=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+        resblock_kernel_sizes_decoder=[3, 7, 11],
+        upsample_rates_decoder=[8, 8, 2, 2],
+        upsample_initial_channel_decoder=512,
+        upsample_kernel_sizes_decoder=[16, 16, 4, 4],
+        d_vector_dim=512,
+        cond_d_vector_in_each_upsampling_layer=True,
+        speaker_encoder_audio_config={
+            "fft_size": 512,
+            "win_length": 400,
+            "hop_length": 160,
+            "sample_rate": 16000,
+            "preemphasis": 0.97,
+            "num_mels": 64,
+        },
+    ):
+        super().__init__()
+        self.input_sample_rate = input_sample_rate
+        self.output_sample_rate = output_sample_rate
+        self.output_hop_length = output_hop_length
+        self.ar_mel_length_compression = ar_mel_length_compression
+        self.speaker_encoder_audio_config = speaker_encoder_audio_config
+        self.waveform_decoder = HifiganGenerator(
+            decoder_input_dim,
+            1,
+            resblock_type_decoder,
+            resblock_dilation_sizes_decoder,
+            resblock_kernel_sizes_decoder,
+            upsample_kernel_sizes_decoder,
+            upsample_initial_channel_decoder,
+            upsample_rates_decoder,
+            inference_padding=0,
+            cond_channels=d_vector_dim,
+            conv_pre_weight_norm=False,
+            conv_post_weight_norm=False,
+            conv_post_bias=False,
+            cond_in_each_up_layer=cond_d_vector_in_each_upsampling_layer,
+        )
+        self.speaker_encoder = ResNetSpeakerEncoder(
+            input_dim=64,
+            proj_dim=512,
+            log_input=True,
+            use_torch_spec=True,
+            audio_config=speaker_encoder_audio_config,
+        )
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def forward(self, latents, g=None):
+        """
+        Args:
+            x (Tensor): feature input tensor (GPT latent).
+            g (Tensor): global conditioning input tensor.
+
+        Returns:
+            Tensor: output waveform.
+
+        Shapes:
+            x: [B, C, T]
+            Tensor: [B, 1, T]
+        """
+
+        z = torch.nn.functional.interpolate(
+            latents.transpose(1, 2),
+            scale_factor=[self.ar_mel_length_compression / self.output_hop_length],
+            mode="linear",
+        ).squeeze(1)
+        # upsample to the right sr
+        if self.output_sample_rate != self.input_sample_rate:
+            z = torch.nn.functional.interpolate(
+                z,
+                scale_factor=[self.output_sample_rate / self.input_sample_rate],
+                mode="linear",
+            ).squeeze(0)
+        o = self.waveform_decoder(z, g=g)
+        return o
+
+    @torch.no_grad()
+    def inference(self, c, g):
+        """
+        Args:
+            x (Tensor): feature input tensor (GPT latent).
+            g (Tensor): global conditioning input tensor.
+
+        Returns:
+            Tensor: output waveform.
+
+        Shapes:
+            x: [B, C, T]
+            Tensor: [B, 1, T]
+        """
+        return self.forward(c, g=g)
+
+    def load_checkpoint(self, checkpoint_path, eval=False):  # pylint: disable=unused-argument, redefined-builtin
+        state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"))
+        # remove unused keys
+        state = state["model"]
+        states_keys = list(state.keys())
+        for key in states_keys:
+            if "waveform_decoder." not in key and "speaker_encoder." not in key:
+                del state[key]
+
+        self.load_state_dict(state)
+        if eval:
+            self.eval()
+            assert not self.training
+            self.waveform_decoder.remove_weight_norm()

TTS/tts/layers/xtts/latent_encoder.py

@@ -0,0 +1,141 @@
+# ported from: Originally ported from: https://github.com/neonbjb/tortoise-tts
+
+import math
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class GroupNorm32(nn.GroupNorm):
+    def forward(self, x):
+        return super().forward(x.float()).type(x.dtype)
+
+
+def conv_nd(dims, *args, **kwargs):
+    if dims == 1:
+        return nn.Conv1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.Conv2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.Conv3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def normalization(channels):
+    groups = 32
+    if channels <= 16:
+        groups = 8
+    elif channels <= 64:
+        groups = 16
+    while channels % groups != 0:
+        groups = int(groups / 2)
+    assert groups > 2
+    return GroupNorm32(groups, channels)
+
+
+def zero_module(module):
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+class QKVAttention(nn.Module):
+    def __init__(self, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+
+    def forward(self, qkv, mask=None, qk_bias=0):
+        """
+        Apply QKV attention.
+
+        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
+        :return: an [N x (H * C) x T] tensor after attention.
+        """
+        bs, width, length = qkv.shape
+        assert width % (3 * self.n_heads) == 0
+        ch = width // (3 * self.n_heads)
+        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
+        scale = 1 / math.sqrt(math.sqrt(ch))
+        weight = torch.einsum("bct,bcs->bts", q * scale, k * scale)  # More stable with f16 than dividing afterwards
+        weight = weight + qk_bias
+        if mask is not None:
+            mask = mask.repeat(self.n_heads, 1, 1)
+            weight[mask.logical_not()] = -torch.inf
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        a = torch.einsum("bts,bcs->bct", weight, v)
+
+        return a.reshape(bs, -1, length)
+
+
+class AttentionBlock(nn.Module):
+    """An attention block that allows spatial positions to attend to each other."""
+
+    def __init__(
+        self,
+        channels,
+        num_heads=1,
+        num_head_channels=-1,
+        out_channels=None,
+        do_activation=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        out_channels = channels if out_channels is None else out_channels
+        self.do_activation = do_activation
+        if num_head_channels == -1:
+            self.num_heads = num_heads
+        else:
+            assert (
+                channels % num_head_channels == 0
+            ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
+            self.num_heads = channels // num_head_channels
+        self.norm = normalization(channels)
+        self.qkv = conv_nd(1, channels, out_channels * 3, 1)
+        self.attention = QKVAttention(self.num_heads)
+
+        self.x_proj = nn.Identity() if out_channels == channels else conv_nd(1, channels, out_channels, 1)
+        self.proj_out = zero_module(conv_nd(1, out_channels, out_channels, 1))
+
+    def forward(self, x, mask=None, qk_bias=0):
+        b, c, *spatial = x.shape
+        if mask is not None:
+            if len(mask.shape) == 2:
+                mask = mask.unsqueeze(0).repeat(x.shape[0], 1, 1)
+            if mask.shape[1] != x.shape[-1]:
+                mask = mask[:, : x.shape[-1], : x.shape[-1]]
+
+        x = x.reshape(b, c, -1)
+        x = self.norm(x)
+        if self.do_activation:
+            x = F.silu(x, inplace=True)
+        qkv = self.qkv(x)
+        h = self.attention(qkv, mask=mask, qk_bias=qk_bias)
+        h = self.proj_out(h)
+        xp = self.x_proj(x)
+        return (xp + h).reshape(b, xp.shape[1], *spatial)
+
+
+class ConditioningEncoder(nn.Module):
+    def __init__(
+        self,
+        spec_dim,
+        embedding_dim,
+        attn_blocks=6,
+        num_attn_heads=4,
+    ):
+        super().__init__()
+        attn = []
+        self.init = nn.Conv1d(spec_dim, embedding_dim, kernel_size=1)
+        for a in range(attn_blocks):
+            attn.append(AttentionBlock(embedding_dim, num_attn_heads))
+        self.attn = nn.Sequential(*attn)
+        self.dim = embedding_dim
+
+    def forward(self, x):
+        """
+        x: (b, 80, s)
+        """
+        h = self.init(x)
+        h = self.attn(h)
+        return h

TTS/tts/layers/xtts/perceiver_encoder.py

@@ -0,0 +1,319 @@
+# Adapted from https://github.com/lucidrains/naturalspeech2-pytorch/blob/659bec7f7543e7747e809e950cc2f84242fbeec7/naturalspeech2_pytorch/naturalspeech2_pytorch.py#L532
+
+from collections import namedtuple
+from functools import wraps
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+from packaging import version
+from torch import einsum, nn
+
+
+def exists(val):
+    return val is not None
+
+
+def once(fn):
+    called = False
+
+    @wraps(fn)
+    def inner(x):
+        nonlocal called
+        if called:
+            return
+        called = True
+        return fn(x)
+
+    return inner
+
+
+print_once = once(print)
+
+# main class
+
+
+class Attend(nn.Module):
+    def __init__(self, dropout=0.0, causal=False, use_flash=False):
+        super().__init__()
+        self.dropout = dropout
+        self.attn_dropout = nn.Dropout(dropout)
+
+        self.causal = causal
+        self.register_buffer("mask", None, persistent=False)
+
+        self.use_flash = use_flash
+        assert not (
+            use_flash and version.parse(torch.__version__) < version.parse("2.0.0")
+        ), "in order to use flash attention, you must be using pytorch 2.0 or above"
+
+        # determine efficient attention configs for cuda and cpu
+        self.config = namedtuple("EfficientAttentionConfig", ["enable_flash", "enable_math", "enable_mem_efficient"])
+        self.cpu_config = self.config(True, True, True)
+        self.cuda_config = None
+
+        if not torch.cuda.is_available() or not use_flash:
+            return
+
+        device_properties = torch.cuda.get_device_properties(torch.device("cuda"))
+
+        if device_properties.major == 8 and device_properties.minor == 0:
+            print_once("A100 GPU detected, using flash attention if input tensor is on cuda")
+            self.cuda_config = self.config(True, False, False)
+        else:
+            print_once("Non-A100 GPU detected, using math or mem efficient attention if input tensor is on cuda")
+            self.cuda_config = self.config(False, True, True)
+
+    def get_mask(self, n, device):
+        if exists(self.mask) and self.mask.shape[-1] >= n:
+            return self.mask[:n, :n]
+
+        mask = torch.ones((n, n), device=device, dtype=torch.bool).triu(1)
+        self.register_buffer("mask", mask, persistent=False)
+        return mask
+
+    def flash_attn(self, q, k, v, mask=None):
+        _, heads, q_len, _, k_len, is_cuda = *q.shape, k.shape[-2], q.is_cuda
+
+        # Recommended for multi-query single-key-value attention by Tri Dao
+        # kv shape torch.Size([1, 512, 64]) -> torch.Size([1, 8, 512, 64])
+
+        if k.ndim == 3:
+            k = rearrange(k, "b ... -> b 1 ...").expand_as(q)
+
+        if v.ndim == 3:
+            v = rearrange(v, "b ... -> b 1 ...").expand_as(q)
+
+        # Check if mask exists and expand to compatible shape
+        # The mask is B L, so it would have to be expanded to B H N L
+
+        if exists(mask):
+            mask = rearrange(mask, "b j -> b 1 1 j")
+            mask = mask.expand(-1, heads, q_len, -1)
+
+        # Check if there is a compatible device for flash attention
+
+        config = self.cuda_config if is_cuda else self.cpu_config
+
+        # pytorch 2.0 flash attn: q, k, v, mask, dropout, causal, softmax_scale
+
+        with torch.backends.cuda.sdp_kernel(**config._asdict()):
+            out = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=mask, dropout_p=self.dropout if self.training else 0.0, is_causal=self.causal
+            )
+
+        return out
+
+    def forward(self, q, k, v, mask=None):
+        """
+        einstein notation
+        b - batch
+        h - heads
+        n, i, j - sequence length (base sequence length, source, target)
+        d - feature dimension
+        """
+
+        n, device = q.shape[-2], q.device
+
+        scale = q.shape[-1] ** -0.5
+
+        if self.use_flash:
+            return self.flash_attn(q, k, v, mask=mask)
+
+        kv_einsum_eq = "b j d" if k.ndim == 3 else "b h j d"
+
+        # similarity
+
+        sim = einsum(f"b h i d, {kv_einsum_eq} -> b h i j", q, k) * scale
+
+        # key padding mask
+
+        if exists(mask):
+            mask = rearrange(mask, "b j -> b 1 1 j")
+            sim = sim.masked_fill(~mask, -torch.finfo(sim.dtype).max)
+
+        # causal mask
+
+        if self.causal:
+            causal_mask = self.get_mask(n, device)
+            sim = sim.masked_fill(causal_mask, -torch.finfo(sim.dtype).max)
+
+        # attention
+
+        attn = sim.softmax(dim=-1)
+        attn = self.attn_dropout(attn)
+
+        # aggregate values
+
+        out = einsum(f"b h i j, {kv_einsum_eq} -> b h i d", attn, v)
+
+        return out
+
+
+def Sequential(*mods):
+    return nn.Sequential(*filter(exists, mods))
+
+
+def exists(x):
+    return x is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if callable(d) else d
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, scale=True, dim_cond=None):
+        super().__init__()
+        self.cond = exists(dim_cond)
+        self.to_gamma_beta = nn.Linear(dim_cond, dim * 2) if self.cond else None
+
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(dim)) if scale else None
+
+    def forward(self, x, cond=None):
+        gamma = default(self.gamma, 1)
+        out = F.normalize(x, dim=-1) * self.scale * gamma
+
+        if not self.cond:
+            return out
+
+        assert exists(cond)
+        gamma, beta = self.to_gamma_beta(cond).chunk(2, dim=-1)
+        gamma, beta = map(lambda t: rearrange(t, "b d -> b 1 d"), (gamma, beta))
+        return out * gamma + beta
+
+
+class CausalConv1d(nn.Conv1d):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        (kernel_size,) = self.kernel_size
+        (dilation,) = self.dilation
+        (stride,) = self.stride
+
+        assert stride == 1
+        self.causal_padding = dilation * (kernel_size - 1)
+
+    def forward(self, x):
+        causal_padded_x = F.pad(x, (self.causal_padding, 0), value=0.0)
+        return super().forward(causal_padded_x)
+
+
+class GEGLU(nn.Module):
+    def forward(self, x):
+        x, gate = x.chunk(2, dim=-1)
+        return F.gelu(gate) * x
+
+
+def FeedForward(dim, mult=4, causal_conv=False):
+    dim_inner = int(dim * mult * 2 / 3)
+
+    conv = None
+    if causal_conv:
+        conv = nn.Sequential(
+            Rearrange("b n d -> b d n"),
+            CausalConv1d(dim_inner, dim_inner, 3),
+            Rearrange("b d n -> b n d"),
+        )
+
+    return Sequential(nn.Linear(dim, dim_inner * 2), GEGLU(), conv, nn.Linear(dim_inner, dim))
+
+
+class PerceiverResampler(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth=2,
+        dim_context=None,
+        num_latents=32,
+        dim_head=64,
+        heads=8,
+        ff_mult=4,
+        use_flash_attn=False,
+    ):
+        super().__init__()
+        dim_context = default(dim_context, dim)
+
+        self.proj_context = nn.Linear(dim_context, dim) if dim_context != dim else nn.Identity()
+
+        self.latents = nn.Parameter(torch.randn(num_latents, dim))
+        nn.init.normal_(self.latents, std=0.02)
+
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        Attention(
+                            dim=dim,
+                            dim_head=dim_head,
+                            heads=heads,
+                            use_flash=use_flash_attn,
+                            cross_attn_include_queries=True,
+                        ),
+                        FeedForward(dim=dim, mult=ff_mult),
+                    ]
+                )
+            )
+
+        self.norm = RMSNorm(dim)
+
+    def forward(self, x, mask=None):
+        batch = x.shape[0]
+
+        x = self.proj_context(x)
+
+        latents = repeat(self.latents, "n d -> b n d", b=batch)
+
+        for attn, ff in self.layers:
+            latents = attn(latents, x, mask=mask) + latents
+            latents = ff(latents) + latents
+
+        return self.norm(latents)
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        *,
+        dim_context=None,
+        causal=False,
+        dim_head=64,
+        heads=8,
+        dropout=0.0,
+        use_flash=False,
+        cross_attn_include_queries=False,
+    ):
+        super().__init__()
+        self.scale = dim_head**-0.5
+        self.heads = heads
+        self.cross_attn_include_queries = cross_attn_include_queries
+
+        dim_inner = dim_head * heads
+        dim_context = default(dim_context, dim)
+
+        self.attend = Attend(causal=causal, dropout=dropout, use_flash=use_flash)
+        self.to_q = nn.Linear(dim, dim_inner, bias=False)
+        self.to_kv = nn.Linear(dim_context, dim_inner * 2, bias=False)
+        self.to_out = nn.Linear(dim_inner, dim, bias=False)
+
+    def forward(self, x, context=None, mask=None):
+        h, has_context = self.heads, exists(context)
+
+        context = default(context, x)
+
+        if has_context and self.cross_attn_include_queries:
+            context = torch.cat((x, context), dim=-2)
+
+        q, k, v = (self.to_q(x), *self.to_kv(context).chunk(2, dim=-1))
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v))
+
+        out = self.attend(q, k, v, mask=mask)
+
+        out = rearrange(out, "b h n d -> b n (h d)")
+        return self.to_out(out)

TTS/tts/layers/xtts/stream_generator.py

@@ -0,0 +1,930 @@
+# Adapted from: https://github.com/LowinLi/transformers-stream-generator
+
+import copy
+import inspect
+import random
+import warnings
+from typing import Callable, List, Optional, Union
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch import nn
+from transformers import (
+    BeamSearchScorer,
+    ConstrainedBeamSearchScorer,
+    DisjunctiveConstraint,
+    GenerationConfig,
+    GenerationMixin,
+    LogitsProcessorList,
+    PhrasalConstraint,
+    PreTrainedModel,
+    StoppingCriteriaList,
+)
+from transformers.generation.utils import GenerateOutput, SampleOutput, logger
+
+
+def setup_seed(seed):
+    if seed == -1:
+        return
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+
+
+class StreamGenerationConfig(GenerationConfig):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.do_stream = kwargs.pop("do_stream", False)
+
+
+class NewGenerationMixin(GenerationMixin):
+    @torch.no_grad()
+    def generate(
+        self,
+        inputs: Optional[torch.Tensor] = None,
+        generation_config: Optional[StreamGenerationConfig] = None,
+        logits_processor: Optional[LogitsProcessorList] = None,
+        stopping_criteria: Optional[StoppingCriteriaList] = None,
+        prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
+        synced_gpus: Optional[bool] = False,
+        seed=0,
+        **kwargs,
+    ) -> Union[GenerateOutput, torch.LongTensor]:
+        r"""
+
+        Generates sequences of token ids for models with a language modeling head.
+
+        <Tip warning={true}>
+
+        Most generation-controlling parameters are set in `generation_config` which, if not passed, will be set to the
+        model's default generation configuration. You can override any `generation_config` by passing the corresponding
+        parameters to generate(), e.g. `.generate(inputs, num_beams=4, do_sample=True)`.
+
+        For an overview of generation strategies and code examples, check out the [following
+        guide](./generation_strategies).
+
+        </Tip>
+
+        Parameters:
+            inputs (`torch.Tensor` of varying shape depending on the modality, *optional*):
+                The sequence used as a prompt for the generation or as model inputs to the encoder. If `None` the
+                method initializes it with `bos_token_id` and a batch size of 1. For decoder-only models `inputs`
+                should of in the format of `input_ids`. For encoder-decoder models *inputs* can represent any of
+                `input_ids`, `input_values`, `input_features`, or `pixel_values`.
+            generation_config (`~generation.GenerationConfig`, *optional*):
+                The generation configuration to be used as base parametrization for the generation call. `**kwargs`
+                passed to generate matching the attributes of `generation_config` will override them. If
+                `generation_config` is not provided, the default will be used, which had the following loading
+                priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
+                configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
+                default values, whose documentation should be checked to parameterize generation.
+            logits_processor (`LogitsProcessorList`, *optional*):
+                Custom logits processors that complement the default logits processors built from arguments and
+                generation config. If a logit processor is passed that is already created with the arguments or a
+                generation config an error is thrown. This feature is intended for advanced users.
+            stopping_criteria (`StoppingCriteriaList`, *optional*):
+                Custom stopping criteria that complement the default stopping criteria built from arguments and a
+                generation config. If a stopping criteria is passed that is already created with the arguments or a
+                generation config an error is thrown. This feature is intended for advanced users.
+            prefix_allowed_tokens_fn (`Callable[[int, torch.Tensor], List[int]]`, *optional*):
+                If provided, this function constraints the beam search to allowed tokens only at each step. If not
+                provided no constraint is applied. This function takes 2 arguments: the batch ID `batch_id` and
+                `input_ids`. It has to return a list with the allowed tokens for the next generation step conditioned
+                on the batch ID `batch_id` and the previously generated tokens `inputs_ids`. This argument is useful
+                for constrained generation conditioned on the prefix, as described in [Autoregressive Entity
+                Retrieval](https://arxiv.org/abs/2010.00904).
+            synced_gpus (`bool`, *optional*, defaults to `False`):
+                Whether to continue running the while loop until max_length (needed for ZeRO stage 3)
+            kwargs:
+                Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be
+                forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder
+                specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*.
+
+        Return:
+            [`~utils.ModelOutput`] or `torch.LongTensor`: A [`~utils.ModelOutput`] (if `return_dict_in_generate=True`
+            or when `config.return_dict_in_generate=True`) or a `torch.FloatTensor`.
+
+                If the model is *not* an encoder-decoder model (`model.config.is_encoder_decoder=False`), the possible
+                [`~utils.ModelOutput`] types are:
+
+                    - [`~generation.GreedySearchDecoderOnlyOutput`],
+                    - [`~generation.SampleDecoderOnlyOutput`],
+                    - [`~generation.BeamSearchDecoderOnlyOutput`],
+                    - [`~generation.BeamSampleDecoderOnlyOutput`]
+
+                If the model is an encoder-decoder model (`model.config.is_encoder_decoder=True`), the possible
+                [`~utils.ModelOutput`] types are:
+
+                    - [`~generation.GreedySearchEncoderDecoderOutput`],
+                    - [`~generation.SampleEncoderDecoderOutput`],
+                    - [`~generation.BeamSearchEncoderDecoderOutput`],
+                    - [`~generation.BeamSampleEncoderDecoderOutput`]
+        """
+        # setup_seed(seed)
+        # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
+        self._validate_model_class()
+
+        # priority: `generation_config` argument > `model.generation_config` (the default generation config)
+        if generation_config is None:
+            # legacy: users may modify the model configuration to control generation -- update the generation config
+            # model attribute accordingly, if it was created from the model config
+            if self.generation_config._from_model_config:
+                new_generation_config = StreamGenerationConfig.from_model_config(self.config)
+                if new_generation_config != self.generation_config:
+                    warnings.warn(
+                        "You have modified the pretrained model configuration to control generation. This is a"
+                        " deprecated strategy to control generation and will be removed soon, in a future version."
+                        " Please use a generation configuration file (see"
+                        " https://huggingface.co/docs/transformers/main_classes/text_generation)"
+                    )
+                    self.generation_config = new_generation_config
+            generation_config = self.generation_config
+
+        generation_config = copy.deepcopy(generation_config)
+        model_kwargs = generation_config.update(**kwargs)  # All unused kwargs must be model kwargs
+        # self._validate_model_kwargs(model_kwargs.copy())
+
+        # 2. Set generation parameters if not already defined
+        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
+        stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
+
+        if generation_config.pad_token_id is None and generation_config.eos_token_id is not None:
+            if model_kwargs.get("attention_mask", None) is None:
+                logger.warning(
+                    "The attention mask and the pad token id were not set. As a consequence, you may observe "
+                    "unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results."
+                )
+            eos_token_id = generation_config.eos_token_id
+            if isinstance(eos_token_id, list):
+                eos_token_id = eos_token_id[0]
+            logger.warning(f"Setting `pad_token_id` to `eos_token_id`:{eos_token_id} for open-end generation.")
+            generation_config.pad_token_id = eos_token_id
+
+        # 3. Define model inputs
+        # inputs_tensor has to be defined
+        # model_input_name is defined if model-specific keyword input is passed
+        # otherwise model_input_name is None
+        # all model-specific keyword inputs are removed from `model_kwargs`
+        inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs(
+            inputs, generation_config.bos_token_id, model_kwargs
+        )
+        batch_size = inputs_tensor.shape[0]
+
+        # 4. Define other model kwargs
+        model_kwargs["output_attentions"] = generation_config.output_attentions
+        model_kwargs["output_hidden_states"] = generation_config.output_hidden_states
+        model_kwargs["use_cache"] = generation_config.use_cache
+
+        accepts_attention_mask = "attention_mask" in set(inspect.signature(self.forward).parameters.keys())
+        requires_attention_mask = "encoder_outputs" not in model_kwargs
+
+        if model_kwargs.get("attention_mask", None) is None and requires_attention_mask and accepts_attention_mask:
+            model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation(
+                inputs_tensor,
+                generation_config.pad_token_id,
+                generation_config.eos_token_id,
+            )
+
+        # decoder-only models should use left-padding for generation
+        if not self.config.is_encoder_decoder:
+            if (
+                generation_config.pad_token_id is not None
+                and torch.sum(inputs_tensor[:, -1] == generation_config.pad_token_id) > 0
+            ):
+                logger.warning(
+                    "A decoder-only architecture is being used, but right-padding was detected! For correct "
+                    "generation results, please set `padding_side='left'` when initializing the tokenizer."
+                )
+
+        if self.config.is_encoder_decoder and "encoder_outputs" not in model_kwargs:
+            # if model is encoder decoder encoder_outputs are created
+            # and added to `model_kwargs`
+            model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(
+                inputs_tensor, model_kwargs, model_input_name
+            )
+
+        # 5. Prepare `input_ids` which will be used for auto-regressive generation
+        if self.config.is_encoder_decoder:
+            input_ids = self._prepare_decoder_input_ids_for_generation(
+                batch_size,
+                decoder_start_token_id=generation_config.decoder_start_token_id,
+                bos_token_id=generation_config.bos_token_id,
+                model_kwargs=model_kwargs,
+                device=inputs_tensor.device,
+            )
+        else:
+            # if decoder-only then inputs_tensor has to be `input_ids`
+            input_ids = inputs_tensor
+
+        # 6. Prepare `max_length` depending on other stopping criteria.
+        input_ids_seq_length = input_ids.shape[-1]
+        has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
+        if has_default_max_length and generation_config.max_new_tokens is None:
+            warnings.warn(
+                "Neither `max_length` nor `max_new_tokens` has been set, `max_length` will default to"
+                f" {generation_config.max_length} (`generation_config.max_length`). Controlling `max_length` via the"
+                " config is deprecated and `max_length` will be removed from the config in v5 of Transformers -- we"
+                " recommend using `max_new_tokens` to control the maximum length of the generation.",
+                UserWarning,
+            )
+        elif has_default_max_length and generation_config.max_new_tokens is not None:
+            generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length
+        elif not has_default_max_length and generation_config.max_new_tokens is not None:
+            raise ValueError(
+                "Both `max_new_tokens` and `max_length` have been set but they serve the same purpose -- setting a"
+                " limit to the generated output length. Remove one of those arguments. Please refer to the"
+                " documentation for more information. "
+                "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)"
+            )
+
+        if generation_config.min_length is not None and generation_config.min_length > generation_config.max_length:
+            raise ValueError(
+                f"Unfeasible length constraints: the minimum length ({generation_config.min_length}) is larger than"
+                f" the maximum length ({generation_config.max_length})"
+            )
+        if input_ids_seq_length >= generation_config.max_length:
+            input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
+            logger.warning(
+                f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"
+                f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
+                " increasing `max_new_tokens`."
+            )
+
+        # 7. determine generation mode
+        is_constraint_gen_mode = (
+            generation_config.constraints is not None or generation_config.force_words_ids is not None
+        )
+
+        is_contrastive_search_gen_mode = (
+            generation_config.top_k is not None
+            and generation_config.top_k > 1
+            and generation_config.do_sample is False
+            and generation_config.penalty_alpha is not None
+            and generation_config.penalty_alpha > 0
+        )
+
+        is_greedy_gen_mode = (
+            (generation_config.num_beams == 1)
+            and (generation_config.num_beam_groups == 1)
+            and generation_config.do_sample is False
+            and not is_constraint_gen_mode
+            and not is_contrastive_search_gen_mode
+        )
+        is_sample_gen_mode = (
+            (generation_config.num_beams == 1)
+            and (generation_config.num_beam_groups == 1)
+            and generation_config.do_sample is True
+            and generation_config.do_stream is False
+            and not is_constraint_gen_mode
+            and not is_contrastive_search_gen_mode
+        )
+        is_sample_gen_stream_mode = (
+            (generation_config.num_beams == 1)
+            and (generation_config.num_beam_groups == 1)
+            and generation_config.do_stream is True
+            and not is_constraint_gen_mode
+            and not is_contrastive_search_gen_mode
+        )
+        is_beam_gen_mode = (
+            (generation_config.num_beams > 1)
+            and (generation_config.num_beam_groups == 1)
+            and generation_config.do_sample is False
+            and not is_constraint_gen_mode
+            and not is_contrastive_search_gen_mode
+        )
+        is_beam_sample_gen_mode = (
+            (generation_config.num_beams > 1)
+            and (generation_config.num_beam_groups == 1)
+            and generation_config.do_sample is True
+            and not is_constraint_gen_mode
+            and not is_contrastive_search_gen_mode
+        )
+        is_group_beam_gen_mode = (
+            (generation_config.num_beams > 1)
+            and (generation_config.num_beam_groups > 1)
+            and not is_constraint_gen_mode
+            and not is_contrastive_search_gen_mode
+        )
+
+        if generation_config.num_beam_groups > generation_config.num_beams:
+            raise ValueError("`num_beam_groups` has to be smaller or equal to `num_beams`")
+        if is_group_beam_gen_mode and generation_config.do_sample is True:
+            raise ValueError(
+                "Diverse beam search cannot be used in sampling mode. Make sure that `do_sample` is set to `False`."
+            )
+
+        if self.device.type != input_ids.device.type:
+            warnings.warn(
+                "You are calling .generate() with the `input_ids` being on a device type different"
+                f" than your model's device. `input_ids` is on {input_ids.device.type}, whereas the model"
+                f" is on {self.device.type}. You may experience unexpected behaviors or slower generation."
+                " Please make sure that you have put `input_ids` to the"
+                f" correct device by calling for example input_ids = input_ids.to('{self.device.type}') before"
+                " running `.generate()`.",
+                UserWarning,
+            )
+        # 8. prepare distribution pre_processing samplers
+        logits_processor = self._get_logits_processor(
+            generation_config=generation_config,
+            input_ids_seq_length=input_ids_seq_length,
+            encoder_input_ids=inputs_tensor,
+            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
+            logits_processor=logits_processor,
+        )
+
+        # 9. prepare stopping criteria
+        stopping_criteria = self._get_stopping_criteria(
+            generation_config=generation_config, stopping_criteria=stopping_criteria
+        )
+        # 10. go into different generation modes
+        if is_greedy_gen_mode:
+            if generation_config.num_return_sequences > 1:
+                raise ValueError(
+                    f"num_return_sequences has to be 1, but is {generation_config.num_return_sequences} when doing"
+                    " greedy search."
+                )
+
+            # 11. run greedy search
+            return self.greedy_search(
+                input_ids,
+                logits_processor=logits_processor,
+                stopping_criteria=stopping_criteria,
+                pad_token_id=generation_config.pad_token_id,
+                eos_token_id=generation_config.eos_token_id,
+                output_scores=generation_config.output_scores,
+                return_dict_in_generate=generation_config.return_dict_in_generate,
+                synced_gpus=synced_gpus,
+                **model_kwargs,
+            )
+
+        elif is_contrastive_search_gen_mode:
+            if generation_config.num_return_sequences > 1:
+                raise ValueError(
+                    f"num_return_sequences has to be 1, but is {generation_config.num_return_sequences} when doing"
+                    " contrastive search."
+                )
+
+            return self.contrastive_search(
+                input_ids,
+                top_k=generation_config.top_k,
+                penalty_alpha=generation_config.penalty_alpha,
+                logits_processor=logits_processor,
+                stopping_criteria=stopping_criteria,
+                pad_token_id=generation_config.pad_token_id,
+                eos_token_id=generation_config.eos_token_id,
+                output_scores=generation_config.output_scores,
+                return_dict_in_generate=generation_config.return_dict_in_generate,
+                synced_gpus=synced_gpus,
+                **model_kwargs,
+            )
+
+        elif is_sample_gen_mode:
+            # 11. prepare logits warper
+            logits_warper = self._get_logits_warper(generation_config)
+
+            # 12. expand input_ids with `num_return_sequences` additional sequences per batch
+            input_ids, model_kwargs = self._expand_inputs_for_generation(
+                input_ids=input_ids,
+                expand_size=generation_config.num_return_sequences,
+                is_encoder_decoder=self.config.is_encoder_decoder,
+                **model_kwargs,
+            )
+
+            # 13. run sample
+            return self.sample(
+                input_ids,
+                logits_processor=logits_processor,
+                logits_warper=logits_warper,
+                stopping_criteria=stopping_criteria,
+                pad_token_id=generation_config.pad_token_id,
+                eos_token_id=generation_config.eos_token_id,
+                output_scores=generation_config.output_scores,
+                return_dict_in_generate=generation_config.return_dict_in_generate,
+                synced_gpus=synced_gpus,
+                **model_kwargs,
+            )
+        elif is_sample_gen_stream_mode:
+            # 11. prepare logits warper
+            logits_warper = self._get_logits_warper(generation_config)
+
+            # 12. expand input_ids with `num_return_sequences` additional sequences per batch
+            input_ids, model_kwargs = self._expand_inputs_for_generation(
+                input_ids=input_ids,
+                expand_size=generation_config.num_return_sequences,
+                is_encoder_decoder=self.config.is_encoder_decoder,
+                **model_kwargs,
+            )
+
+            # 13. run sample
+            return self.sample_stream(
+                input_ids,
+                logits_processor=logits_processor,
+                logits_warper=logits_warper,
+                stopping_criteria=stopping_criteria,
+                pad_token_id=generation_config.pad_token_id,
+                eos_token_id=generation_config.eos_token_id,
+                output_scores=generation_config.output_scores,
+                return_dict_in_generate=generation_config.return_dict_in_generate,
+                synced_gpus=synced_gpus,
+                **model_kwargs,
+            )
+        elif is_beam_gen_mode:
+            if generation_config.num_return_sequences > generation_config.num_beams:
+                raise ValueError("`num_return_sequences` has to be smaller or equal to `num_beams`.")
+
+            if stopping_criteria.max_length is None:
+                raise ValueError("`max_length` needs to be a stopping_criteria for now.")
+
+            # 11. prepare beam search scorer
+            beam_scorer = BeamSearchScorer(
+                batch_size=batch_size,
+                num_beams=generation_config.num_beams,
+                device=inputs_tensor.device,
+                length_penalty=generation_config.length_penalty,
+                do_early_stopping=generation_config.early_stopping,
+                num_beam_hyps_to_keep=generation_config.num_return_sequences,
+            )
+            # 12. interleave input_ids with `num_beams` additional sequences per batch
+            input_ids, model_kwargs = self._expand_inputs_for_generation(
+                input_ids=input_ids,
+                expand_size=generation_config.num_beams,
+                is_encoder_decoder=self.config.is_encoder_decoder,
+                **model_kwargs,
+            )
+            # 13. run beam search
+            return self.beam_search(
+                input_ids,
+                beam_scorer,
+                logits_processor=logits_processor,
+                stopping_criteria=stopping_criteria,
+                pad_token_id=generation_config.pad_token_id,
+                eos_token_id=generation_config.eos_token_id,
+                output_scores=generation_config.output_scores,
+                return_dict_in_generate=generation_config.return_dict_in_generate,
+                synced_gpus=synced_gpus,
+                **model_kwargs,
+            )
+
+        elif is_beam_sample_gen_mode:
+            # 11. prepare logits warper
+            logits_warper = self._get_logits_warper(generation_config)
+
+            if stopping_criteria.max_length is None:
+                raise ValueError("`max_length` needs to be a stopping_criteria for now.")
+            # 12. prepare beam search scorer
+            beam_scorer = BeamSearchScorer(
+                batch_size=batch_size * generation_config.num_return_sequences,
+                num_beams=generation_config.num_beams,
+                device=inputs_tensor.device,
+                length_penalty=generation_config.length_penalty,
+                do_early_stopping=generation_config.early_stopping,
+            )
+
+            # 13. interleave input_ids with `num_beams` additional sequences per batch
+            input_ids, model_kwargs = self._expand_inputs_for_generation(
+                input_ids=input_ids,
+                expand_size=generation_config.num_beams * generation_config.num_return_sequences,
+                is_encoder_decoder=self.config.is_encoder_decoder,
+                **model_kwargs,
+            )
+
+            # 14. run beam sample
+            return self.beam_sample(
+                input_ids,
+                beam_scorer,
+                logits_processor=logits_processor,
+                logits_warper=logits_warper,
+                stopping_criteria=stopping_criteria,
+                pad_token_id=generation_config.pad_token_id,
+                eos_token_id=generation_config.eos_token_id,
+                output_scores=generation_config.output_scores,
+                return_dict_in_generate=generation_config.return_dict_in_generate,
+                synced_gpus=synced_gpus,
+                **model_kwargs,
+            )
+
+        elif is_group_beam_gen_mode:
+            if generation_config.num_return_sequences > generation_config.num_beams:
+                raise ValueError("`num_return_sequences` has to be smaller or equal to `num_beams`.")
+
+            if generation_config.num_beams % generation_config.num_beam_groups != 0:
+                raise ValueError("`num_beams` should be divisible by `num_beam_groups` for group beam search.")
+
+            if stopping_criteria.max_length is None:
+                raise ValueError("`max_length` needs to be a stopping_criteria for now.")
+
+            has_default_typical_p = kwargs.get("typical_p") is None and generation_config.typical_p == 1.0
+            if not has_default_typical_p:
+                raise ValueError("Decoder argument `typical_p` is not supported with beam groups.")
+
+            # 11. prepare beam search scorer
+            beam_scorer = BeamSearchScorer(
+                batch_size=batch_size,
+                num_beams=generation_config.num_beams,
+                max_length=stopping_criteria.max_length,
+                device=inputs_tensor.device,
+                length_penalty=generation_config.length_penalty,
+                do_early_stopping=generation_config.early_stopping,
+                num_beam_hyps_to_keep=generation_config.num_return_sequences,
+                num_beam_groups=generation_config.num_beam_groups,
+            )
+            # 12. interleave input_ids with `num_beams` additional sequences per batch
+            input_ids, model_kwargs = self._expand_inputs_for_generation(
+                input_ids=input_ids,
+                expand_size=generation_config.num_beams,
+                is_encoder_decoder=self.config.is_encoder_decoder,
+                **model_kwargs,
+            )
+            # 13. run beam search
+            return self.group_beam_search(
+                input_ids,
+                beam_scorer,
+                logits_processor=logits_processor,
+                stopping_criteria=stopping_criteria,
+                pad_token_id=generation_config.pad_token_id,
+                eos_token_id=generation_config.eos_token_id,
+                output_scores=generation_config.output_scores,
+                return_dict_in_generate=generation_config.return_dict_in_generate,
+                synced_gpus=synced_gpus,
+                **model_kwargs,
+            )
+
+        elif is_constraint_gen_mode:
+            if generation_config.num_return_sequences > generation_config.num_beams:
+                raise ValueError("`num_return_sequences` has to be smaller or equal to `num_beams`.")
+
+            if stopping_criteria.max_length is None:
+                raise ValueError("`max_length` needs to be a stopping_criteria for now.")
+
+            if generation_config.num_beams <= 1:
+                raise ValueError("`num_beams` needs to be greater than 1 for constrained generation.")
+
+            if generation_config.do_sample:
+                raise ValueError("`do_sample` needs to be false for constrained generation.")
+
+            if generation_config.num_beam_groups is not None and generation_config.num_beam_groups > 1:
+                raise ValueError("`num_beam_groups` not supported yet for constrained generation.")
+
+            final_constraints = []
+            if generation_config.constraints is not None:
+                final_constraints = generation_config.constraints
+
+            if generation_config.force_words_ids is not None:
+
+                def typeerror():
+                    raise ValueError(
+                        "`force_words_ids` has to either be a `List[List[List[int]]]` or `List[List[int]]`"
+                        f"of positive integers, but is {generation_config.force_words_ids}."
+                    )
+
+                if (
+                    not isinstance(generation_config.force_words_ids, list)
+                    or len(generation_config.force_words_ids) == 0
+                ):
+                    typeerror()
+
+                for word_ids in generation_config.force_words_ids:
+                    if isinstance(word_ids[0], list):
+                        if not isinstance(word_ids, list) or len(word_ids) == 0:
+                            typeerror()
+                        if any(not isinstance(token_ids, list) for token_ids in word_ids):
+                            typeerror()
+                        if any(
+                            any((not isinstance(token_id, int) or token_id < 0) for token_id in token_ids)
+                            for token_ids in word_ids
+                        ):
+                            typeerror()
+
+                        constraint = DisjunctiveConstraint(word_ids)
+                    else:
+                        if not isinstance(word_ids, list) or len(word_ids) == 0:
+                            typeerror()
+                        if any((not isinstance(token_id, int) or token_id < 0) for token_id in word_ids):
+                            typeerror()
+
+                        constraint = PhrasalConstraint(word_ids)
+                    final_constraints.append(constraint)
+
+            # 11. prepare beam search scorer
+            constrained_beam_scorer = ConstrainedBeamSearchScorer(
+                constraints=final_constraints,
+                batch_size=batch_size,
+                num_beams=generation_config.num_beams,
+                device=inputs_tensor.device,
+                length_penalty=generation_config.length_penalty,
+                do_early_stopping=generation_config.early_stopping,
+                num_beam_hyps_to_keep=generation_config.num_return_sequences,
+            )
+            # 12. interleave input_ids with `num_beams` additional sequences per batch
+            input_ids, model_kwargs = self._expand_inputs_for_generation(
+                input_ids=input_ids,
+                expand_size=generation_config.num_beams,
+                is_encoder_decoder=self.config.is_encoder_decoder,
+                **model_kwargs,
+            )
+            # 13. run beam search
+            return self.constrained_beam_search(
+                input_ids,
+                constrained_beam_scorer=constrained_beam_scorer,
+                logits_processor=logits_processor,
+                stopping_criteria=stopping_criteria,
+                pad_token_id=generation_config.pad_token_id,
+                eos_token_id=generation_config.eos_token_id,
+                output_scores=generation_config.output_scores,
+                return_dict_in_generate=generation_config.return_dict_in_generate,
+                synced_gpus=synced_gpus,
+                **model_kwargs,
+            )
+
+    @torch.no_grad()
+    def sample_stream(
+        self,
+        input_ids: torch.LongTensor,
+        logits_processor: Optional[LogitsProcessorList] = None,
+        stopping_criteria: Optional[StoppingCriteriaList] = None,
+        logits_warper: Optional[LogitsProcessorList] = None,
+        max_length: Optional[int] = None,
+        pad_token_id: Optional[int] = None,
+        eos_token_id: Optional[Union[int, List[int]]] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_scores: Optional[bool] = None,
+        return_dict_in_generate: Optional[bool] = None,
+        synced_gpus: Optional[bool] = False,
+        **model_kwargs,
+    ) -> Union[SampleOutput, torch.LongTensor]:
+        r"""
+        Generates sequences of token ids for models with a language modeling head using **multinomial sampling** and
+        can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models.
+
+        <Tip warning={true}>
+
+        In most cases, you do not need to call [`~generation.GenerationMixin.sample`] directly. Use generate() instead.
+        For an overview of generation strategies and code examples, check the [following
+        guide](./generation_strategies).
+
+        </Tip>
+
+        Parameters:
+            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                The sequence used as a prompt for the generation.
+            logits_processor (`LogitsProcessorList`, *optional*):
+                An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`]
+                used to modify the prediction scores of the language modeling head applied at each generation step.
+            stopping_criteria (`StoppingCriteriaList`, *optional*):
+                An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`]
+                used to tell if the generation loop should stop.
+            logits_warper (`LogitsProcessorList`, *optional*):
+                An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsWarper`] used
+                to warp the prediction score distribution of the language modeling head applied before multinomial
+                sampling at each generation step.
+            max_length (`int`, *optional*, defaults to 20):
+                **DEPRECATED**. Use `logits_processor` or `stopping_criteria` directly to cap the number of generated
+                tokens. The maximum length of the sequence to be generated.
+            pad_token_id (`int`, *optional*):
+                The id of the *padding* token.
+            eos_token_id (`int`, *optional*):
+                The id of the *end-of-sequence* token.
+            output_attentions (`bool`, *optional*, defaults to `False`):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more details.
+            output_hidden_states (`bool`, *optional*, defaults to `False`):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
+                for more details.
+            output_scores (`bool`, *optional*, defaults to `False`):
+                Whether or not to return the prediction scores. See `scores` under returned tensors for more details.
+            return_dict_in_generate (`bool`, *optional*, defaults to `False`):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+            synced_gpus (`bool`, *optional*, defaults to `False`):
+                Whether to continue running the while loop until max_length (needed for ZeRO stage 3)
+            model_kwargs:
+                Additional model specific kwargs will be forwarded to the `forward` function of the model. If model is
+                an encoder-decoder model the kwargs should include `encoder_outputs`.
+
+        Return:
+            [`~generation.SampleDecoderOnlyOutput`], [`~generation.SampleEncoderDecoderOutput`] or `torch.LongTensor`:
+            A `torch.LongTensor` containing the generated tokens (default behaviour) or a
+            [`~generation.SampleDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and
+            `return_dict_in_generate=True` or a [`~generation.SampleEncoderDecoderOutput`] if
+            `model.config.is_encoder_decoder=True`.
+
+        Examples:
+
+        ```python
+        >>> from transformers import (
+        ...     AutoTokenizer,
+        ...     AutoModelForCausalLM,
+        ...     LogitsProcessorList,
+        ...     MinLengthLogitsProcessor,
+        ...     TopKLogitsWarper,
+        ...     TemperatureLogitsWarper,
+        ...     StoppingCriteriaList,
+        ...     MaxLengthCriteria,
+        ... )
+        >>> import torch
+
+        >>> tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        >>> model = AutoModelForCausalLM.from_pretrained("gpt2")
+
+        >>> # set pad_token_id to eos_token_id because GPT2 does not have a EOS token
+        >>> model.config.pad_token_id = model.config.eos_token_id
+        >>> model.generation_config.pad_token_id = model.config.eos_token_id
+
+        >>> input_prompt = "Today is a beautiful day, and"
+        >>> input_ids = tokenizer(input_prompt, return_tensors="pt").input_ids
+
+        >>> # instantiate logits processors
+        >>> logits_processor = LogitsProcessorList(
+        ...     [
+        ...         MinLengthLogitsProcessor(15, eos_token_id=model.generation_config.eos_token_id),
+        ...     ]
+        ... )
+        >>> # instantiate logits processors
+        >>> logits_warper = LogitsProcessorList(
+        ...     [
+        ...         TopKLogitsWarper(50),
+        ...         TemperatureLogitsWarper(0.7),
+        ...     ]
+        ... )
+
+        >>> stopping_criteria = StoppingCriteriaList([MaxLengthCriteria(max_length=20)])
+
+        >>> torch.manual_seed(0)  # doctest: +IGNORE_RESULT
+        >>> outputs = model.sample(
+        ...     input_ids,
+        ...     logits_processor=logits_processor,
+        ...     logits_warper=logits_warper,
+        ...     stopping_criteria=stopping_criteria,
+        ... )
+
+        >>> tokenizer.batch_decode(outputs, skip_special_tokens=True)
+        ['Today is a beautiful day, and a wonderful day.\n\nI was lucky enough to meet the']
+        ```"""
+        # init values
+        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
+        stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
+        if max_length is not None:
+            warnings.warn(
+                "`max_length` is deprecated in this function, use"
+                " `stopping_criteria=StoppingCriteriaList(MaxLengthCriteria(max_length=max_length))` instead.",
+                UserWarning,
+            )
+            stopping_criteria = validate_stopping_criteria(stopping_criteria, max_length)
+        logits_warper = logits_warper if logits_warper is not None else LogitsProcessorList()
+        pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id
+        eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id
+        if isinstance(eos_token_id, int):
+            eos_token_id = [eos_token_id]
+        output_scores = output_scores if output_scores is not None else self.generation_config.output_scores
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.generation_config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.generation_config.output_hidden_states
+        )
+        return_dict_in_generate = (
+            return_dict_in_generate
+            if return_dict_in_generate is not None
+            else self.generation_config.return_dict_in_generate
+        )
+
+        # init attention / hidden states / scores tuples
+        scores = () if (return_dict_in_generate and output_scores) else None
+        decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+        cross_attentions = () if (return_dict_in_generate and output_attentions) else None
+        decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+
+        # keep track of which sequences are already finished
+        unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)
+
+        this_peer_finished = False  # used by synced_gpus only
+        # auto-regressive generation
+        while True:
+            if synced_gpus:
+                # Under synced_gpus the `forward` call must continue until all gpus complete their sequence.
+                # The following logic allows an early break if all peers finished generating their sequence
+                this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0).to(input_ids.device)
+                # send 0.0 if we finished, 1.0 otherwise
+                dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM)
+                # did all peers finish? the reduced sum will be 0.0 then
+                if this_peer_finished_flag.item() == 0.0:
+                    break
+
+            # prepare model inputs
+            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+
+            # forward pass to get next token
+            outputs = self(
+                **model_inputs,
+                return_dict=True,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+            )
+
+            if synced_gpus and this_peer_finished:
+                continue  # don't waste resources running the code we don't need
+
+            next_token_logits = outputs.logits[:, -1, :]
+
+            # pre-process distribution
+            next_token_scores = logits_processor(input_ids, next_token_logits)
+            next_token_scores = logits_warper(input_ids, next_token_scores)
+
+            # Store scores, attentions and hidden_states when required
+            if return_dict_in_generate:
+                if output_scores:
+                    scores += (next_token_scores,)
+                if output_attentions:
+                    decoder_attentions += (
+                        (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,)
+                    )
+                    if self.config.is_encoder_decoder:
+                        cross_attentions += (outputs.cross_attentions,)
+
+                if output_hidden_states:
+                    decoder_hidden_states += (
+                        (outputs.decoder_hidden_states,) if self.config.is_encoder_decoder else (outputs.hidden_states,)
+                    )
+
+            # sample
+            probs = nn.functional.softmax(next_token_scores, dim=-1)
+            next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+
+            # finished sentences should have their next token be a padding token
+            if eos_token_id is not None:
+                if pad_token_id is None:
+                    raise ValueError("If `eos_token_id` is defined, make sure that `pad_token_id` is defined.")
+                next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)
+            yield next_tokens, self.final_norm(outputs.hidden_states[-1][:, -1])
+            # update generated ids, model inputs, and length for next step
+            input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
+            model_kwargs = self._update_model_kwargs_for_generation(
+                outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
+            )
+
+            # if eos_token was found in one sentence, set sentence to finished
+            if eos_token_id is not None:
+                unfinished_sequences = unfinished_sequences.mul((sum(next_tokens != i for i in eos_token_id)).long())
+
+            # stop when each sentence is finished, or if we exceed the maximum length
+            if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores):
+                if not synced_gpus:
+                    break
+                else:
+                    this_peer_finished = True
+
+
+def init_stream_support():
+    """Overload PreTrainedModel for streaming."""
+    PreTrainedModel.generate_stream = NewGenerationMixin.generate
+    PreTrainedModel.sample_stream = NewGenerationMixin.sample_stream
+
+
+if __name__ == "__main__":
+    from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel
+
+    PreTrainedModel.generate = NewGenerationMixin.generate
+    PreTrainedModel.sample_stream = NewGenerationMixin.sample_stream
+    model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m", torch_dtype=torch.float16)
+
+    tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-560m")
+    model = model.to("cuda:0")
+    model = model.eval()
+    prompt_text = "hello? \n"
+    input_ids = tokenizer(prompt_text, return_tensors="pt", add_special_tokens=False).input_ids
+    input_ids = input_ids.to("cuda:0")
+
+    with torch.no_grad():
+        result = model.generate(
+            input_ids,
+            max_new_tokens=200,
+            do_sample=True,
+            top_k=30,
+            top_p=0.85,
+            temperature=0.35,
+            repetition_penalty=1.2,
+            early_stopping=True,
+            seed=0,
+        )
+        print(tokenizer.decode(result, skip_special_tokens=True))
+        generator = model.generate(
+            input_ids,
+            max_new_tokens=200,
+            do_sample=True,
+            top_k=30,
+            top_p=0.85,
+            temperature=0.35,
+            repetition_penalty=1.2,
+            early_stopping=True,
+            seed=0,
+            do_stream=True,
+        )
+        stream_result = ""
+        for x in generator:
+            chunk = tokenizer.decode(x, skip_special_tokens=True)
+            stream_result += chunk
+        print(stream_result)

TTS/tts/layers/xtts/tokenizer.py

@@ -0,0 +1,843 @@
+import os
+import re
+import textwrap
+from functools import cached_property
+
+import pypinyin
+import torch
+from hangul_romanize import Transliter
+from hangul_romanize.rule import academic
+from num2words import num2words
+from spacy.lang.ar import Arabic
+from spacy.lang.en import English
+from spacy.lang.es import Spanish
+from spacy.lang.ja import Japanese
+from spacy.lang.zh import Chinese
+from tokenizers import Tokenizer
+
+from TTS.tts.layers.xtts.zh_num2words import TextNorm as zh_num2words
+
+
+def get_spacy_lang(lang):
+    if lang == "zh":
+        return Chinese()
+    elif lang == "ja":
+        return Japanese()
+    elif lang == "ar":
+        return Arabic()
+    elif lang == "es":
+        return Spanish()
+    else:
+        # For most languages, Enlish does the job
+        return English()
+
+
+def split_sentence(text, lang, text_split_length=250):
+    """Preprocess the input text"""
+    text_splits = []
+    if text_split_length is not None and len(text) >= text_split_length:
+        text_splits.append("")
+        nlp = get_spacy_lang(lang)
+        nlp.add_pipe("sentencizer")
+        doc = nlp(text)
+        for sentence in doc.sents:
+            if len(text_splits[-1]) + len(str(sentence)) <= text_split_length:
+                # if the last sentence + the current sentence is less than the text_split_length
+                # then add the current sentence to the last sentence
+                text_splits[-1] += " " + str(sentence)
+                text_splits[-1] = text_splits[-1].lstrip()
+            elif len(str(sentence)) > text_split_length:
+                # if the current sentence is greater than the text_split_length
+                for line in textwrap.wrap(
+                    str(sentence),
+                    width=text_split_length,
+                    drop_whitespace=True,
+                    break_on_hyphens=False,
+                    tabsize=1,
+                ):
+                    text_splits.append(str(line))
+            else:
+                text_splits.append(str(sentence))
+
+        if len(text_splits) > 1:
+            if text_splits[0] == "":
+                del text_splits[0]
+    else:
+        text_splits = [text.lstrip()]
+
+    return text_splits
+
+
+_whitespace_re = re.compile(r"\s+")
+
+# List of (regular expression, replacement) pairs for abbreviations:
+_abbreviations = {
+    "en": [
+        (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+        for x in [
+            ("mrs", "misess"),
+            ("mr", "mister"),
+            ("dr", "doctor"),
+            ("st", "saint"),
+            ("co", "company"),
+            ("jr", "junior"),
+            ("maj", "major"),
+            ("gen", "general"),
+            ("drs", "doctors"),
+            ("rev", "reverend"),
+            ("lt", "lieutenant"),
+            ("hon", "honorable"),
+            ("sgt", "sergeant"),
+            ("capt", "captain"),
+            ("esq", "esquire"),
+            ("ltd", "limited"),
+            ("col", "colonel"),
+            ("ft", "fort"),
+        ]
+    ],
+    "es": [
+        (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+        for x in [
+            ("sra", "señora"),
+            ("sr", "señor"),
+            ("dr", "doctor"),
+            ("dra", "doctora"),
+            ("st", "santo"),
+            ("co", "compañía"),
+            ("jr", "junior"),
+            ("ltd", "limitada"),
+        ]
+    ],
+    "fr": [
+        (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+        for x in [
+            ("mme", "madame"),
+            ("mr", "monsieur"),
+            ("dr", "docteur"),
+            ("st", "saint"),
+            ("co", "compagnie"),
+            ("jr", "junior"),
+            ("ltd", "limitée"),
+        ]
+    ],
+    "de": [
+        (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+        for x in [
+            ("fr", "frau"),
+            ("dr", "doktor"),
+            ("st", "sankt"),
+            ("co", "firma"),
+            ("jr", "junior"),
+        ]
+    ],
+    "pt": [
+        (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+        for x in [
+            ("sra", "senhora"),
+            ("sr", "senhor"),
+            ("dr", "doutor"),
+            ("dra", "doutora"),
+            ("st", "santo"),
+            ("co", "companhia"),
+            ("jr", "júnior"),
+            ("ltd", "limitada"),
+        ]
+    ],
+    "it": [
+        (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+        for x in [
+            # ("sig.ra", "signora"),
+            ("sig", "signore"),
+            ("dr", "dottore"),
+            ("st", "santo"),
+            ("co", "compagnia"),
+            ("jr", "junior"),
+            ("ltd", "limitata"),
+        ]
+    ],
+    "pl": [
+        (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+        for x in [
+            ("p", "pani"),
+            ("m", "pan"),
+            ("dr", "doktor"),
+            ("sw", "święty"),
+            ("jr", "junior"),
+        ]
+    ],
+    "ar": [
+        (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+        for x in [
+            # There are not many common abbreviations in Arabic as in English.
+        ]
+    ],
+    "zh": [
+        (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+        for x in [
+            # Chinese doesn't typically use abbreviations in the same way as Latin-based scripts.
+        ]
+    ],
+    "cs": [
+        (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+        for x in [
+            ("dr", "doktor"),  # doctor
+            ("ing", "inženýr"),  # engineer
+            ("p", "pan"),  # Could also map to pani for woman but no easy way to do it
+            # Other abbreviations would be specialized and not as common.
+        ]
+    ],
+    "ru": [
+        (re.compile("\\b%s\\b" % x[0], re.IGNORECASE), x[1])
+        for x in [
+            ("г-жа", "госпожа"),  # Mrs.
+            ("г-н", "господин"),  # Mr.
+            ("д-р", "доктор"),  # doctor
+            # Other abbreviations are less common or specialized.
+        ]
+    ],
+    "nl": [
+        (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+        for x in [
+            ("dhr", "de heer"),  # Mr.
+            ("mevr", "mevrouw"),  # Mrs.
+            ("dr", "dokter"),  # doctor
+            ("jhr", "jonkheer"),  # young lord or nobleman
+            # Dutch uses more abbreviations, but these are the most common ones.
+        ]
+    ],
+    "tr": [
+        (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+        for x in [
+            ("b", "bay"),  # Mr.
+            ("byk", "büyük"),  # büyük
+            ("dr", "doktor"),  # doctor
+            # Add other Turkish abbreviations here if needed.
+        ]
+    ],
+    "hu": [
+        (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+        for x in [
+            ("dr", "doktor"),  # doctor
+            ("b", "bácsi"),  # Mr.
+            ("nőv", "nővér"),  # nurse
+            # Add other Hungarian abbreviations here if needed.
+        ]
+    ],
+    "ko": [
+        (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+        for x in [
+            # Korean doesn't typically use abbreviations in the same way as Latin-based scripts.
+        ]
+    ],
+}
+
+
+def expand_abbreviations_multilingual(text, lang="en"):
+    for regex, replacement in _abbreviations[lang]:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+_symbols_multilingual = {
+    "en": [
+        (re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
+        for x in [
+            ("&", " and "),
+            ("@", " at "),
+            ("%", " percent "),
+            ("#", " hash "),
+            ("$", " dollar "),
+            ("£", " pound "),
+            ("°", " degree "),
+        ]
+    ],
+    "es": [
+        (re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
+        for x in [
+            ("&", " y "),
+            ("@", " arroba "),
+            ("%", " por ciento "),
+            ("#", " numeral "),
+            ("$", " dolar "),
+            ("£", " libra "),
+            ("°", " grados "),
+        ]
+    ],
+    "fr": [
+        (re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
+        for x in [
+            ("&", " et "),
+            ("@", " arobase "),
+            ("%", " pour cent "),
+            ("#", " dièse "),
+            ("$", " dollar "),
+            ("£", " livre "),
+            ("°", " degrés "),
+        ]
+    ],
+    "de": [
+        (re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
+        for x in [
+            ("&", " und "),
+            ("@", " at "),
+            ("%", " prozent "),
+            ("#", " raute "),
+            ("$", " dollar "),
+            ("£", " pfund "),
+            ("°", " grad "),
+        ]
+    ],
+    "pt": [
+        (re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
+        for x in [
+            ("&", " e "),
+            ("@", " arroba "),
+            ("%", " por cento "),
+            ("#", " cardinal "),
+            ("$", " dólar "),
+            ("£", " libra "),
+            ("°", " graus "),
+        ]
+    ],
+    "it": [
+        (re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
+        for x in [
+            ("&", " e "),
+            ("@", " chiocciola "),
+            ("%", " per cento "),
+            ("#", " cancelletto "),
+            ("$", " dollaro "),
+            ("£", " sterlina "),
+            ("°", " gradi "),
+        ]
+    ],
+    "pl": [
+        (re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
+        for x in [
+            ("&", " i "),
+            ("@", " małpa "),
+            ("%", " procent "),
+            ("#", " krzyżyk "),
+            ("$", " dolar "),
+            ("£", " funt "),
+            ("°", " stopnie "),
+        ]
+    ],
+    "ar": [
+        # Arabic
+        (re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
+        for x in [
+            ("&", " و "),
+            ("@", " على "),
+            ("%", " في المئة "),
+            ("#", " رقم "),
+            ("$", " دولار "),
+            ("£", " جنيه "),
+            ("°", " درجة "),
+        ]
+    ],
+    "zh": [
+        # Chinese
+        (re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
+        for x in [
+            ("&", " 和 "),
+            ("@", " 在 "),
+            ("%", " 百分之 "),
+            ("#", " 号 "),
+            ("$", " 美元 "),
+            ("£", " 英镑 "),
+            ("°", " 度 "),
+        ]
+    ],
+    "cs": [
+        # Czech
+        (re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
+        for x in [
+            ("&", " a "),
+            ("@", " na "),
+            ("%", " procento "),
+            ("#", " křížek "),
+            ("$", " dolar "),
+            ("£", " libra "),
+            ("°", " stupně "),
+        ]
+    ],
+    "ru": [
+        # Russian
+        (re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
+        for x in [
+            ("&", " и "),
+            ("@", " собака "),
+            ("%", " процентов "),
+            ("#", " номер "),
+            ("$", " доллар "),
+            ("£", " фунт "),
+            ("°", " градус "),
+        ]
+    ],
+    "nl": [
+        # Dutch
+        (re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
+        for x in [
+            ("&", " en "),
+            ("@", " bij "),
+            ("%", " procent "),
+            ("#", " hekje "),
+            ("$", " dollar "),
+            ("£", " pond "),
+            ("°", " graden "),
+        ]
+    ],
+    "tr": [
+        (re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
+        for x in [
+            ("&", " ve "),
+            ("@", " at "),
+            ("%", " yüzde "),
+            ("#", " diyez "),
+            ("$", " dolar "),
+            ("£", " sterlin "),
+            ("°", " derece "),
+        ]
+    ],
+    "hu": [
+        (re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
+        for x in [
+            ("&", " és "),
+            ("@", " kukac "),
+            ("%", " százalék "),
+            ("#", " kettőskereszt "),
+            ("$", " dollár "),
+            ("£", " font "),
+            ("°", " fok "),
+        ]
+    ],
+    "ko": [
+        # Korean
+        (re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
+        for x in [
+            ("&", " 그리고 "),
+            ("@", " 에 "),
+            ("%", " 퍼센트 "),
+            ("#", " 번호 "),
+            ("$", " 달러 "),
+            ("£", " 파운드 "),
+            ("°", " 도 "),
+        ]
+    ],
+}
+
+
+def expand_symbols_multilingual(text, lang="en"):
+    for regex, replacement in _symbols_multilingual[lang]:
+        text = re.sub(regex, replacement, text)
+        text = text.replace("  ", " ")  # Ensure there are no double spaces
+    return text.strip()
+
+
+_ordinal_re = {
+    "en": re.compile(r"([0-9]+)(st|nd|rd|th)"),
+    "es": re.compile(r"([0-9]+)(º|ª|er|o|a|os|as)"),
+    "fr": re.compile(r"([0-9]+)(º|ª|er|re|e|ème)"),
+    "de": re.compile(r"([0-9]+)(st|nd|rd|th|º|ª|\.(?=\s|$))"),
+    "pt": re.compile(r"([0-9]+)(º|ª|o|a|os|as)"),
+    "it": re.compile(r"([0-9]+)(º|°|ª|o|a|i|e)"),
+    "pl": re.compile(r"([0-9]+)(º|ª|st|nd|rd|th)"),
+    "ar": re.compile(r"([0-9]+)(ون|ين|ث|ر|ى)"),
+    "cs": re.compile(r"([0-9]+)\.(?=\s|$)"),  # In Czech, a dot is often used after the number to indicate ordinals.
+    "ru": re.compile(r"([0-9]+)(-й|-я|-е|-ое|-ье|-го)"),
+    "nl": re.compile(r"([0-9]+)(de|ste|e)"),
+    "tr": re.compile(r"([0-9]+)(\.|inci|nci|uncu|üncü|\.)"),
+    "hu": re.compile(r"([0-9]+)(\.|adik|edik|odik|edik|ödik|ödike|ik)"),
+    "ko": re.compile(r"([0-9]+)(번째|번|차|째)"),
+}
+_number_re = re.compile(r"[0-9]+")
+_currency_re = {
+    "USD": re.compile(r"((\$[0-9\.\,]*[0-9]+)|([0-9\.\,]*[0-9]+\$))"),
+    "GBP": re.compile(r"((£[0-9\.\,]*[0-9]+)|([0-9\.\,]*[0-9]+£))"),
+    "EUR": re.compile(r"(([0-9\.\,]*[0-9]+€)|((€[0-9\.\,]*[0-9]+)))"),
+}
+
+_comma_number_re = re.compile(r"\b\d{1,3}(,\d{3})*(\.\d+)?\b")
+_dot_number_re = re.compile(r"\b\d{1,3}(.\d{3})*(\,\d+)?\b")
+_decimal_number_re = re.compile(r"([0-9]+[.,][0-9]+)")
+
+
+def _remove_commas(m):
+    text = m.group(0)
+    if "," in text:
+        text = text.replace(",", "")
+    return text
+
+
+def _remove_dots(m):
+    text = m.group(0)
+    if "." in text:
+        text = text.replace(".", "")
+    return text
+
+
+def _expand_decimal_point(m, lang="en"):
+    amount = m.group(1).replace(",", ".")
+    return num2words(float(amount), lang=lang if lang != "cs" else "cz")
+
+
+def _expand_currency(m, lang="en", currency="USD"):
+    amount = float((re.sub(r"[^\d.]", "", m.group(0).replace(",", "."))))
+    full_amount = num2words(amount, to="currency", currency=currency, lang=lang if lang != "cs" else "cz")
+
+    and_equivalents = {
+        "en": ", ",
+        "es": " con ",
+        "fr": " et ",
+        "de": " und ",
+        "pt": " e ",
+        "it": " e ",
+        "pl": ", ",
+        "cs": ", ",
+        "ru": ", ",
+        "nl": ", ",
+        "ar": ", ",
+        "tr": ", ",
+        "hu": ", ",
+        "ko": ", ",
+    }
+
+    if amount.is_integer():
+        last_and = full_amount.rfind(and_equivalents[lang])
+        if last_and != -1:
+            full_amount = full_amount[:last_and]
+
+    return full_amount
+
+
+def _expand_ordinal(m, lang="en"):
+    return num2words(int(m.group(1)), ordinal=True, lang=lang if lang != "cs" else "cz")
+
+
+def _expand_number(m, lang="en"):
+    return num2words(int(m.group(0)), lang=lang if lang != "cs" else "cz")
+
+
+def expand_numbers_multilingual(text, lang="en"):
+    if lang == "zh":
+        text = zh_num2words()(text)
+    else:
+        if lang in ["en", "ru"]:
+            text = re.sub(_comma_number_re, _remove_commas, text)
+        else:
+            text = re.sub(_dot_number_re, _remove_dots, text)
+        try:
+            text = re.sub(_currency_re["GBP"], lambda m: _expand_currency(m, lang, "GBP"), text)
+            text = re.sub(_currency_re["USD"], lambda m: _expand_currency(m, lang, "USD"), text)
+            text = re.sub(_currency_re["EUR"], lambda m: _expand_currency(m, lang, "EUR"), text)
+        except:
+            pass
+        if lang != "tr":
+            text = re.sub(_decimal_number_re, lambda m: _expand_decimal_point(m, lang), text)
+        text = re.sub(_ordinal_re[lang], lambda m: _expand_ordinal(m, lang), text)
+        text = re.sub(_number_re, lambda m: _expand_number(m, lang), text)
+    return text
+
+
+def lowercase(text):
+    return text.lower()
+
+
+def collapse_whitespace(text):
+    return re.sub(_whitespace_re, " ", text)
+
+
+def multilingual_cleaners(text, lang):
+    text = text.replace('"', "")
+    if lang == "tr":
+        text = text.replace("İ", "i")
+        text = text.replace("Ö", "ö")
+        text = text.replace("Ü", "ü")
+    text = lowercase(text)
+    text = expand_numbers_multilingual(text, lang)
+    text = expand_abbreviations_multilingual(text, lang)
+    text = expand_symbols_multilingual(text, lang=lang)
+    text = collapse_whitespace(text)
+    return text
+
+
+def basic_cleaners(text):
+    """Basic pipeline that lowercases and collapses whitespace without transliteration."""
+    text = lowercase(text)
+    text = collapse_whitespace(text)
+    return text
+
+
+def chinese_transliterate(text):
+    return "".join(
+        [p[0] for p in pypinyin.pinyin(text, style=pypinyin.Style.TONE3, heteronym=False, neutral_tone_with_five=True)]
+    )
+
+
+def japanese_cleaners(text, katsu):
+    text = katsu.romaji(text)
+    text = lowercase(text)
+    return text
+
+
+def korean_transliterate(text):
+    r = Transliter(academic)
+    return r.translit(text)
+
+
+DEFAULT_VOCAB_FILE = os.path.join(os.path.dirname(os.path.realpath(__file__)), "../data/tokenizer.json")
+
+
+class VoiceBpeTokenizer:
+    def __init__(self, vocab_file=None):
+        self.tokenizer = None
+        if vocab_file is not None:
+            self.tokenizer = Tokenizer.from_file(vocab_file)
+        self.char_limits = {
+            "en": 250,
+            "de": 253,
+            "fr": 273,
+            "es": 239,
+            "it": 213,
+            "pt": 203,
+            "pl": 224,
+            "zh": 82,
+            "ar": 166,
+            "cs": 186,
+            "ru": 182,
+            "nl": 251,
+            "tr": 226,
+            "ja": 71,
+            "hu": 224,
+            "ko": 95,
+        }
+
+    @cached_property
+    def katsu(self):
+        import cutlet
+
+        return cutlet.Cutlet()
+
+    def check_input_length(self, txt, lang):
+        lang = lang.split("-")[0]  # remove the region
+        limit = self.char_limits.get(lang, 250)
+        if len(txt) > limit:
+            print(
+                f"[!] Warning: The text length exceeds the character limit of {limit} for language '{lang}', this might cause truncated audio."
+            )
+
+    def preprocess_text(self, txt, lang):
+        if lang in {"ar", "cs", "de", "en", "es", "fr", "hu", "it", "nl", "pl", "pt", "ru", "tr", "zh", "ko"}:
+            txt = multilingual_cleaners(txt, lang)
+            if lang == "zh":
+                txt = chinese_transliterate(txt)
+            if lang == "ko":
+                txt = korean_transliterate(txt)
+        elif lang == "ja":
+            txt = japanese_cleaners(txt, self.katsu)
+        elif lang == "hi":
+            # @manmay will implement this
+            txt = basic_cleaners(txt)
+        else:
+            raise NotImplementedError(f"Language '{lang}' is not supported.")
+        return txt
+
+    def encode(self, txt, lang):
+        lang = lang.split("-")[0]  # remove the region
+        self.check_input_length(txt, lang)
+        txt = self.preprocess_text(txt, lang)
+        lang = "zh-cn" if lang == "zh" else lang
+        txt = f"[{lang}]{txt}"
+        txt = txt.replace(" ", "[SPACE]")
+        return self.tokenizer.encode(txt).ids
+
+    def decode(self, seq):
+        if isinstance(seq, torch.Tensor):
+            seq = seq.cpu().numpy()
+        txt = self.tokenizer.decode(seq, skip_special_tokens=False).replace(" ", "")
+        txt = txt.replace("[SPACE]", " ")
+        txt = txt.replace("[STOP]", "")
+        txt = txt.replace("[UNK]", "")
+        return txt
+
+    def __len__(self):
+        return self.tokenizer.get_vocab_size()
+
+    def get_number_tokens(self):
+        return max(self.tokenizer.get_vocab().values()) + 1
+
+
+def test_expand_numbers_multilingual():
+    test_cases = [
+        # English
+        ("In 12.5 seconds.", "In twelve point five seconds.", "en"),
+        ("There were 50 soldiers.", "There were fifty soldiers.", "en"),
+        ("This is a 1st test", "This is a first test", "en"),
+        ("That will be $20 sir.", "That will be twenty dollars sir.", "en"),
+        ("That will be 20€ sir.", "That will be twenty euro sir.", "en"),
+        ("That will be 20.15€ sir.", "That will be twenty euro, fifteen cents sir.", "en"),
+        ("That's 100,000.5.", "That's one hundred thousand point five.", "en"),
+        # French
+        ("En 12,5 secondes.", "En douze virgule cinq secondes.", "fr"),
+        ("Il y avait 50 soldats.", "Il y avait cinquante soldats.", "fr"),
+        ("Ceci est un 1er test", "Ceci est un premier test", "fr"),
+        ("Cela vous fera $20 monsieur.", "Cela vous fera vingt dollars monsieur.", "fr"),
+        ("Cela vous fera 20€ monsieur.", "Cela vous fera vingt euros monsieur.", "fr"),
+        ("Cela vous fera 20,15€ monsieur.", "Cela vous fera vingt euros et quinze centimes monsieur.", "fr"),
+        ("Ce sera 100.000,5.", "Ce sera cent mille virgule cinq.", "fr"),
+        # German
+        ("In 12,5 Sekunden.", "In zwölf Komma fünf Sekunden.", "de"),
+        ("Es gab 50 Soldaten.", "Es gab fünfzig Soldaten.", "de"),
+        ("Dies ist ein 1. Test", "Dies ist ein erste Test", "de"),  # Issue with gender
+        ("Das macht $20 Herr.", "Das macht zwanzig Dollar Herr.", "de"),
+        ("Das macht 20€ Herr.", "Das macht zwanzig Euro Herr.", "de"),
+        ("Das macht 20,15€ Herr.", "Das macht zwanzig Euro und fünfzehn Cent Herr.", "de"),
+        # Spanish
+        ("En 12,5 segundos.", "En doce punto cinco segundos.", "es"),
+        ("Había 50 soldados.", "Había cincuenta soldados.", "es"),
+        ("Este es un 1er test", "Este es un primero test", "es"),
+        ("Eso le costará $20 señor.", "Eso le costará veinte dólares señor.", "es"),
+        ("Eso le costará 20€ señor.", "Eso le costará veinte euros señor.", "es"),
+        ("Eso le costará 20,15€ señor.", "Eso le costará veinte euros con quince céntimos señor.", "es"),
+        # Italian
+        ("In 12,5 secondi.", "In dodici virgola cinque secondi.", "it"),
+        ("C'erano 50 soldati.", "C'erano cinquanta soldati.", "it"),
+        ("Questo è un 1° test", "Questo è un primo test", "it"),
+        ("Ti costerà $20 signore.", "Ti costerà venti dollari signore.", "it"),
+        ("Ti costerà 20€ signore.", "Ti costerà venti euro signore.", "it"),
+        ("Ti costerà 20,15€ signore.", "Ti costerà venti euro e quindici centesimi signore.", "it"),
+        # Portuguese
+        ("Em 12,5 segundos.", "Em doze vírgula cinco segundos.", "pt"),
+        ("Havia 50 soldados.", "Havia cinquenta soldados.", "pt"),
+        ("Este é um 1º teste", "Este é um primeiro teste", "pt"),
+        ("Isso custará $20 senhor.", "Isso custará vinte dólares senhor.", "pt"),
+        ("Isso custará 20€ senhor.", "Isso custará vinte euros senhor.", "pt"),
+        (
+            "Isso custará 20,15€ senhor.",
+            "Isso custará vinte euros e quinze cêntimos senhor.",
+            "pt",
+        ),  # "cêntimos" should be "centavos" num2words issue
+        # Polish
+        ("W 12,5 sekundy.", "W dwanaście przecinek pięć sekundy.", "pl"),
+        ("Było 50 żołnierzy.", "Było pięćdziesiąt żołnierzy.", "pl"),
+        ("To będzie kosztować 20€ panie.", "To będzie kosztować dwadzieścia euro panie.", "pl"),
+        ("To będzie kosztować 20,15€ panie.", "To będzie kosztować dwadzieścia euro, piętnaście centów panie.", "pl"),
+        # Arabic
+        ("في الـ 12,5 ثانية.", "في الـ اثنا عشر  , خمسون ثانية.", "ar"),
+        ("كان هناك 50 جنديًا.", "كان هناك خمسون جنديًا.", "ar"),
+        # ("ستكون النتيجة $20 يا سيد.", 'ستكون النتيجة عشرون دولار يا سيد.', 'ar'), # $ and € are mising from num2words
+        # ("ستكون النتيجة 20€ يا سيد.", 'ستكون النتيجة عشرون يورو يا سيد.', 'ar'),
+        # Czech
+        ("Za 12,5 vteřiny.", "Za dvanáct celá pět vteřiny.", "cs"),
+        ("Bylo tam 50 vojáků.", "Bylo tam padesát vojáků.", "cs"),
+        ("To bude stát 20€ pane.", "To bude stát dvacet euro pane.", "cs"),
+        ("To bude 20.15€ pane.", "To bude dvacet euro, patnáct centů pane.", "cs"),
+        # Russian
+        ("Через 12.5 секунды.", "Через двенадцать запятая пять секунды.", "ru"),
+        ("Там было 50 солдат.", "Там было пятьдесят солдат.", "ru"),
+        ("Это будет 20.15€ сэр.", "Это будет двадцать евро, пятнадцать центов сэр.", "ru"),
+        ("Это будет стоить 20€ господин.", "Это будет стоить двадцать евро господин.", "ru"),
+        # Dutch
+        ("In 12,5 seconden.", "In twaalf komma vijf seconden.", "nl"),
+        ("Er waren 50 soldaten.", "Er waren vijftig soldaten.", "nl"),
+        ("Dat wordt dan $20 meneer.", "Dat wordt dan twintig dollar meneer.", "nl"),
+        ("Dat wordt dan 20€ meneer.", "Dat wordt dan twintig euro meneer.", "nl"),
+        # Chinese (Simplified)
+        ("在12.5秒内", "在十二点五秒内", "zh"),
+        ("有50名士兵", "有五十名士兵", "zh"),
+        # ("那将是$20先生", '那将是二十美元先生', 'zh'), currency doesn't work
+        # ("那将是20€先生", '那将是二十欧元先生', 'zh'),
+        # Turkish
+        # ("12,5 saniye içinde.", 'On iki virgül beş saniye içinde.', 'tr'), # decimal doesn't work for TR
+        ("50 asker vardı.", "elli asker vardı.", "tr"),
+        ("Bu 1. test", "Bu birinci test", "tr"),
+        # ("Bu 100.000,5.", 'Bu yüz bin virgül beş.', 'tr'),
+        # Hungarian
+        ("12,5 másodperc alatt.", "tizenkettő egész öt tized másodperc alatt.", "hu"),
+        ("50 katona volt.", "ötven katona volt.", "hu"),
+        ("Ez az 1. teszt", "Ez az első teszt", "hu"),
+        # Korean
+        ("12.5 초 안에.", "십이 점 다섯 초 안에.", "ko"),
+        ("50 명의 병사가 있었다.", "오십 명의 병사가 있었다.", "ko"),
+        ("이것은 1 번째 테스트입니다", "이것은 첫 번째 테스트입니다", "ko"),
+    ]
+    for a, b, lang in test_cases:
+        out = expand_numbers_multilingual(a, lang=lang)
+        assert out == b, f"'{out}' vs '{b}'"
+
+
+def test_abbreviations_multilingual():
+    test_cases = [
+        # English
+        ("Hello Mr. Smith.", "Hello mister Smith.", "en"),
+        ("Dr. Jones is here.", "doctor Jones is here.", "en"),
+        # Spanish
+        ("Hola Sr. Garcia.", "Hola señor Garcia.", "es"),
+        ("La Dra. Martinez es muy buena.", "La doctora Martinez es muy buena.", "es"),
+        # French
+        ("Bonjour Mr. Dupond.", "Bonjour monsieur Dupond.", "fr"),
+        ("Mme. Moreau est absente aujourd'hui.", "madame Moreau est absente aujourd'hui.", "fr"),
+        # German
+        ("Frau Dr. Müller ist sehr klug.", "Frau doktor Müller ist sehr klug.", "de"),
+        # Portuguese
+        ("Olá Sr. Silva.", "Olá senhor Silva.", "pt"),
+        ("Dra. Costa, você está disponível?", "doutora Costa, você está disponível?", "pt"),
+        # Italian
+        ("Buongiorno, Sig. Rossi.", "Buongiorno, signore Rossi.", "it"),
+        # ("Sig.ra Bianchi, posso aiutarti?", 'signora Bianchi, posso aiutarti?', 'it'), # Issue with matching that pattern
+        # Polish
+        ("Dzień dobry, P. Kowalski.", "Dzień dobry, pani Kowalski.", "pl"),
+        ("M. Nowak, czy mogę zadać pytanie?", "pan Nowak, czy mogę zadać pytanie?", "pl"),
+        # Czech
+        ("P. Novák", "pan Novák", "cs"),
+        ("Dr. Vojtěch", "doktor Vojtěch", "cs"),
+        # Dutch
+        ("Dhr. Jansen", "de heer Jansen", "nl"),
+        ("Mevr. de Vries", "mevrouw de Vries", "nl"),
+        # Russian
+        ("Здравствуйте Г-н Иванов.", "Здравствуйте господин Иванов.", "ru"),
+        ("Д-р Смирнов здесь, чтобы увидеть вас.", "доктор Смирнов здесь, чтобы увидеть вас.", "ru"),
+        # Turkish
+        ("Merhaba B. Yılmaz.", "Merhaba bay Yılmaz.", "tr"),
+        ("Dr. Ayşe burada.", "doktor Ayşe burada.", "tr"),
+        # Hungarian
+        ("Dr. Szabó itt van.", "doktor Szabó itt van.", "hu"),
+    ]
+
+    for a, b, lang in test_cases:
+        out = expand_abbreviations_multilingual(a, lang=lang)
+        assert out == b, f"'{out}' vs '{b}'"
+
+
+def test_symbols_multilingual():
+    test_cases = [
+        ("I have 14% battery", "I have 14 percent battery", "en"),
+        ("Te veo @ la fiesta", "Te veo arroba la fiesta", "es"),
+        ("J'ai 14° de fièvre", "J'ai 14 degrés de fièvre", "fr"),
+        ("Die Rechnung beträgt £ 20", "Die Rechnung beträgt pfund 20", "de"),
+        ("O meu email é ana&joao@gmail.com", "O meu email é ana e joao arroba gmail.com", "pt"),
+        ("linguaggio di programmazione C#", "linguaggio di programmazione C cancelletto", "it"),
+        ("Moja temperatura to 36.6°", "Moja temperatura to 36.6 stopnie", "pl"),
+        ("Mám 14% baterie", "Mám 14 procento baterie", "cs"),
+        ("Těším se na tebe @ party", "Těším se na tebe na party", "cs"),
+        ("У меня 14% заряда", "У меня 14 процентов заряда", "ru"),
+        ("Я буду @ дома", "Я буду собака дома", "ru"),
+        ("Ik heb 14% batterij", "Ik heb 14 procent batterij", "nl"),
+        ("Ik zie je @ het feest", "Ik zie je bij het feest", "nl"),
+        ("لدي 14% في البطارية", "لدي 14 في المئة في البطارية", "ar"),
+        ("我的电量为 14%", "我的电量为 14 百分之", "zh"),
+        ("Pilim %14 dolu.", "Pilim yüzde 14 dolu.", "tr"),
+        ("Az akkumulátorom töltöttsége 14%", "Az akkumulátorom töltöttsége 14 százalék", "hu"),
+        ("배터리 잔량이 14%입니다.", "배터리 잔량이 14 퍼센트입니다.", "ko"),
+    ]
+
+    for a, b, lang in test_cases:
+        out = expand_symbols_multilingual(a, lang=lang)
+        assert out == b, f"'{out}' vs '{b}'"
+
+
+if __name__ == "__main__":
+    test_expand_numbers_multilingual()
+    test_abbreviations_multilingual()
+    test_symbols_multilingual()

TTS/tts/layers/xtts/trainer/dataset.py

@@ -0,0 +1,239 @@
+import os
+import random
+import sys
+
+import torch
+import torch.nn.functional as F
+import torch.utils.data
+
+from TTS.tts.models.xtts import load_audio
+
+torch.set_num_threads(1)
+
+
+def key_samples_by_col(samples, col):
+    """Returns a dictionary of samples keyed by language."""
+    samples_by_col = {}
+    for sample in samples:
+        col_val = sample[col]
+        assert isinstance(col_val, str)
+        if col_val not in samples_by_col:
+            samples_by_col[col_val] = []
+        samples_by_col[col_val].append(sample)
+    return samples_by_col
+
+
+def get_prompt_slice(gt_path, max_sample_length, min_sample_length, sample_rate, is_eval=False):
+    rel_clip = load_audio(gt_path, sample_rate)
+    # if eval uses a middle size sample when it is possible to be more reproducible
+    if is_eval:
+        sample_length = int((min_sample_length + max_sample_length) / 2)
+    else:
+        sample_length = random.randint(min_sample_length, max_sample_length)
+    gap = rel_clip.shape[-1] - sample_length
+    if gap < 0:
+        sample_length = rel_clip.shape[-1] // 2
+    gap = rel_clip.shape[-1] - sample_length
+
+    # if eval start always from the position 0 to be more reproducible
+    if is_eval:
+        rand_start = 0
+    else:
+        rand_start = random.randint(0, gap)
+
+    rand_end = rand_start + sample_length
+    rel_clip = rel_clip[:, rand_start:rand_end]
+    rel_clip = F.pad(rel_clip, pad=(0, max_sample_length - rel_clip.shape[-1]))
+    cond_idxs = [rand_start, rand_end]
+    return rel_clip, rel_clip.shape[-1], cond_idxs
+
+
+class XTTSDataset(torch.utils.data.Dataset):
+    def __init__(self, config, samples, tokenizer, sample_rate, is_eval=False):
+        self.config = config
+        model_args = config.model_args
+        self.failed_samples = set()
+        self.debug_failures = model_args.debug_loading_failures
+        self.max_conditioning_length = model_args.max_conditioning_length
+        self.min_conditioning_length = model_args.min_conditioning_length
+        self.is_eval = is_eval
+        self.tokenizer = tokenizer
+        self.sample_rate = sample_rate
+        self.max_wav_len = model_args.max_wav_length
+        self.max_text_len = model_args.max_text_length
+        self.use_masking_gt_prompt_approach = model_args.gpt_use_masking_gt_prompt_approach
+        assert self.max_wav_len is not None and self.max_text_len is not None
+
+        self.samples = samples
+        if not is_eval:
+            random.seed(config.training_seed)
+            # random.shuffle(self.samples)
+            random.shuffle(self.samples)
+            # order by language
+            self.samples = key_samples_by_col(self.samples, "language")
+            print(" > Sampling by language:", self.samples.keys())
+        else:
+            # for evaluation load and check samples that are corrupted to ensures the reproducibility
+            self.check_eval_samples()
+
+    def check_eval_samples(self):
+        print(" > Filtering invalid eval samples!!")
+        new_samples = []
+        for sample in self.samples:
+            try:
+                tseq, _, wav, _, _, _ = self.load_item(sample)
+            except:
+                continue
+            # Basically, this audio file is nonexistent or too long to be supported by the dataset.
+            if (
+                wav is None
+                or (self.max_wav_len is not None and wav.shape[-1] > self.max_wav_len)
+                or (self.max_text_len is not None and tseq.shape[0] > self.max_text_len)
+            ):
+                continue
+            new_samples.append(sample)
+        self.samples = new_samples
+        print(" > Total eval samples after filtering:", len(self.samples))
+
+    def get_text(self, text, lang):
+        tokens = self.tokenizer.encode(text, lang)
+        tokens = torch.IntTensor(tokens)
+        assert not torch.any(tokens == 1), f"UNK token found in {text} -> {self.tokenizer.decode(tokens)}"
+        # The stop token should always be sacred.
+        assert not torch.any(tokens == 0), f"Stop token found in {text}"
+        return tokens
+
+    def load_item(self, sample):
+        text = str(sample["text"])
+        tseq = self.get_text(text, sample["language"])
+        audiopath = sample["audio_file"]
+        wav = load_audio(audiopath, self.sample_rate)
+        if text is None or len(text.strip()) == 0:
+            raise ValueError
+        if wav is None or wav.shape[-1] < (0.5 * self.sample_rate):
+            # Ultra short clips are also useless (and can cause problems within some models).
+            raise ValueError
+
+        if self.use_masking_gt_prompt_approach:
+            # get a slice from GT to condition the model
+            cond, _, cond_idxs = get_prompt_slice(
+                audiopath, self.max_conditioning_length, self.min_conditioning_length, self.sample_rate, self.is_eval
+            )
+            # if use masking do not use cond_len
+            cond_len = torch.nan
+        else:
+            ref_sample = (
+                sample["reference_path"]
+                if "reference_path" in sample and sample["reference_path"] is not None
+                else audiopath
+            )
+            cond, cond_len, _ = get_prompt_slice(
+                ref_sample, self.max_conditioning_length, self.min_conditioning_length, self.sample_rate, self.is_eval
+            )
+            # if do not use masking use cond_len
+            cond_idxs = torch.nan
+
+        return tseq, audiopath, wav, cond, cond_len, cond_idxs
+
+    def __getitem__(self, index):
+        if self.is_eval:
+            sample = self.samples[index]
+            sample_id = str(index)
+        else:
+            # select a random language
+            lang = random.choice(list(self.samples.keys()))
+            # select random sample
+            index = random.randint(0, len(self.samples[lang]) - 1)
+            sample = self.samples[lang][index]
+            # a unique id for each sampel to deal with fails
+            sample_id = lang + "_" + str(index)
+
+        # ignore samples that we already know that is not valid ones
+        if sample_id in self.failed_samples:
+            if self.debug_failures:
+                print(f"Ignoring sample {sample['audio_file']} because it was already ignored before !!")
+            # call get item again to get other sample
+            return self[1]
+
+        # try to load the sample, if fails added it to the failed samples list
+        try:
+            tseq, audiopath, wav, cond, cond_len, cond_idxs = self.load_item(sample)
+        except:
+            if self.debug_failures:
+                print(f"error loading {sample['audio_file']} {sys.exc_info()}")
+            self.failed_samples.add(sample_id)
+            return self[1]
+
+        # check if the audio and text size limits and if it out of the limits, added it failed_samples
+        if (
+            wav is None
+            or (self.max_wav_len is not None and wav.shape[-1] > self.max_wav_len)
+            or (self.max_text_len is not None and tseq.shape[0] > self.max_text_len)
+        ):
+            # Basically, this audio file is nonexistent or too long to be supported by the dataset.
+            # It's hard to handle this situation properly. Best bet is to return the a random valid token and skew the dataset somewhat as a result.
+            if self.debug_failures and wav is not None and tseq is not None:
+                print(
+                    f"error loading {sample['audio_file']}: ranges are out of bounds; {wav.shape[-1]}, {tseq.shape[0]}"
+                )
+            self.failed_samples.add(sample_id)
+            return self[1]
+
+        res = {
+            # 'real_text': text,
+            "text": tseq,
+            "text_lengths": torch.tensor(tseq.shape[0], dtype=torch.long),
+            "wav": wav,
+            "wav_lengths": torch.tensor(wav.shape[-1], dtype=torch.long),
+            "filenames": audiopath,
+            "conditioning": cond.unsqueeze(1),
+            "cond_lens": torch.tensor(cond_len, dtype=torch.long)
+            if cond_len is not torch.nan
+            else torch.tensor([cond_len]),
+            "cond_idxs": torch.tensor(cond_idxs) if cond_idxs is not torch.nan else torch.tensor([cond_idxs]),
+        }
+        return res
+
+    def __len__(self):
+        if self.is_eval:
+            return len(self.samples)
+        return sum([len(v) for v in self.samples.values()])
+
+    def collate_fn(self, batch):
+        # convert list of dicts to dict of lists
+        B = len(batch)
+
+        batch = {k: [dic[k] for dic in batch] for k in batch[0]}
+
+        # stack for features that already have the same shape
+        batch["wav_lengths"] = torch.stack(batch["wav_lengths"])
+        batch["text_lengths"] = torch.stack(batch["text_lengths"])
+        batch["conditioning"] = torch.stack(batch["conditioning"])
+        batch["cond_lens"] = torch.stack(batch["cond_lens"])
+        batch["cond_idxs"] = torch.stack(batch["cond_idxs"])
+
+        if torch.any(batch["cond_idxs"].isnan()):
+            batch["cond_idxs"] = None
+
+        if torch.any(batch["cond_lens"].isnan()):
+            batch["cond_lens"] = None
+
+        max_text_len = batch["text_lengths"].max()
+        max_wav_len = batch["wav_lengths"].max()
+
+        # create padding tensors
+        text_padded = torch.IntTensor(B, max_text_len)
+        wav_padded = torch.FloatTensor(B, 1, max_wav_len)
+
+        # initialize tensors for zero padding
+        text_padded = text_padded.zero_()
+        wav_padded = wav_padded.zero_()
+        for i in range(B):
+            text = batch["text"][i]
+            text_padded[i, : batch["text_lengths"][i]] = torch.IntTensor(text)
+            wav = batch["wav"][i]
+            wav_padded[i, :, : batch["wav_lengths"][i]] = torch.FloatTensor(wav)
+
+        batch["wav"] = wav_padded
+        batch["padded_text"] = text_padded
+        return batch

TTS/tts/layers/xtts/trainer/gpt_trainer.py

@@ -0,0 +1,504 @@
+from dataclasses import dataclass, field
+from typing import Dict, List, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torchaudio
+from coqpit import Coqpit
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from trainer.torch import DistributedSampler
+from trainer.trainer_utils import get_optimizer, get_scheduler
+
+from TTS.tts.configs.xtts_config import XttsConfig
+from TTS.tts.datasets.dataset import TTSDataset
+from TTS.tts.layers.tortoise.arch_utils import TorchMelSpectrogram
+from TTS.tts.layers.xtts.dvae import DiscreteVAE
+from TTS.tts.layers.xtts.tokenizer import VoiceBpeTokenizer
+from TTS.tts.layers.xtts.trainer.dataset import XTTSDataset
+from TTS.tts.models.base_tts import BaseTTS
+from TTS.tts.models.xtts import Xtts, XttsArgs, XttsAudioConfig
+from TTS.utils.io import load_fsspec
+
+
+@dataclass
+class GPTTrainerConfig(XttsConfig):
+    lr: float = 5e-06
+    training_seed: int = 1
+    optimizer_wd_only_on_weights: bool = False
+    weighted_loss_attrs: dict = field(default_factory=lambda: {})
+    weighted_loss_multipliers: dict = field(default_factory=lambda: {})
+    test_sentences: List[dict] = field(default_factory=lambda: [])
+
+
+@dataclass
+class XttsAudioConfig(XttsAudioConfig):
+    dvae_sample_rate: int = 22050
+
+
+@dataclass
+class GPTArgs(XttsArgs):
+    min_conditioning_length: int = 66150
+    max_conditioning_length: int = 132300
+    gpt_loss_text_ce_weight: float = 0.01
+    gpt_loss_mel_ce_weight: float = 1.0
+    gpt_num_audio_tokens: int = 8194
+    debug_loading_failures: bool = False
+    max_wav_length: int = 255995  # ~11.6 seconds
+    max_text_length: int = 200
+    tokenizer_file: str = ""
+    mel_norm_file: str = "https://coqui.gateway.scarf.sh/v0.14.0_models/mel_norms.pth"
+    dvae_checkpoint: str = ""
+    xtts_checkpoint: str = ""
+    gpt_checkpoint: str = ""  # if defined it will replace the gpt weights on xtts model
+    vocoder: str = ""  # overide vocoder key on the config to avoid json write issues
+
+
+def callback_clearml_load_save(operation_type, model_info):
+    # return None means skip the file upload/log, returning model_info will continue with the log/upload
+    # you can also change the upload destination file name model_info.upload_filename or check the local file size with Path(model_info.local_model_path).stat().st_size
+    assert operation_type in ("load", "save")
+    # print(operation_type, model_info.__dict__)
+
+    if "similarities.pth" in model_info.__dict__["local_model_path"]:
+        return None
+
+    return model_info
+
+
+class GPTTrainer(BaseTTS):
+    def __init__(self, config: Coqpit):
+        """
+        Tortoise GPT training class
+        """
+        super().__init__(config, ap=None, tokenizer=None)
+        self.config = config
+        # init XTTS model
+        self.xtts = Xtts(self.config)
+        # create the tokenizer with the target vocabulary
+        self.xtts.tokenizer = VoiceBpeTokenizer(self.args.tokenizer_file)
+        # init gpt encoder and hifigan decoder
+        self.xtts.init_models()
+
+        if self.args.xtts_checkpoint:
+            self.load_checkpoint(self.config, self.args.xtts_checkpoint, eval=False, strict=False)
+
+        # set mel stats
+        if self.args.mel_norm_file:
+            self.xtts.mel_stats = load_fsspec(self.args.mel_norm_file)
+
+        # load GPT if available
+        if self.args.gpt_checkpoint:
+            gpt_checkpoint = torch.load(self.args.gpt_checkpoint, map_location=torch.device("cpu"))
+            # deal with coqui Trainer exported model
+            if "model" in gpt_checkpoint.keys() and "config" in gpt_checkpoint.keys():
+                print("Coqui Trainer checkpoint detected! Converting it!")
+                gpt_checkpoint = gpt_checkpoint["model"]
+                states_keys = list(gpt_checkpoint.keys())
+                for key in states_keys:
+                    if "gpt." in key:
+                        new_key = key.replace("gpt.", "")
+                        gpt_checkpoint[new_key] = gpt_checkpoint[key]
+                        del gpt_checkpoint[key]
+                    else:
+                        del gpt_checkpoint[key]
+
+            # edit checkpoint if the number of tokens is changed to ensures the better transfer learning possible
+            if (
+                "text_embedding.weight" in gpt_checkpoint
+                and gpt_checkpoint["text_embedding.weight"].shape != self.xtts.gpt.text_embedding.weight.shape
+            ):
+                num_new_tokens = (
+                    self.xtts.gpt.text_embedding.weight.shape[0] - gpt_checkpoint["text_embedding.weight"].shape[0]
+                )
+                print(f" > Loading checkpoint with {num_new_tokens} additional tokens.")
+
+                # add new tokens to a linear layer (text_head)
+                emb_g = gpt_checkpoint["text_embedding.weight"]
+                new_row = torch.randn(num_new_tokens, emb_g.shape[1])
+                start_token_row = emb_g[-1, :]
+                emb_g = torch.cat([emb_g, new_row], axis=0)
+                emb_g[-1, :] = start_token_row
+                gpt_checkpoint["text_embedding.weight"] = emb_g
+
+                # add new weights to the linear layer (text_head)
+                text_head_weight = gpt_checkpoint["text_head.weight"]
+                start_token_row = text_head_weight[-1, :]
+                new_entry = torch.randn(num_new_tokens, self.xtts.gpt.text_head.weight.shape[1])
+                text_head_weight = torch.cat([text_head_weight, new_entry], axis=0)
+                text_head_weight[-1, :] = start_token_row
+                gpt_checkpoint["text_head.weight"] = text_head_weight
+
+                # add new biases to the linear layer (text_head)
+                text_head_bias = gpt_checkpoint["text_head.bias"]
+                start_token_row = text_head_bias[-1]
+                new_bias_entry = torch.zeros(num_new_tokens)
+                text_head_bias = torch.cat([text_head_bias, new_bias_entry], axis=0)
+                text_head_bias[-1] = start_token_row
+                gpt_checkpoint["text_head.bias"] = text_head_bias
+
+            self.xtts.gpt.load_state_dict(gpt_checkpoint, strict=True)
+            print(">> GPT weights restored from:", self.args.gpt_checkpoint)
+
+        # Mel spectrogram extractor for conditioning
+        if self.args.gpt_use_perceiver_resampler:
+            self.torch_mel_spectrogram_style_encoder = TorchMelSpectrogram(
+                filter_length=2048,
+                hop_length=256,
+                win_length=1024,
+                normalize=False,
+                sampling_rate=config.audio.sample_rate,
+                mel_fmin=0,
+                mel_fmax=8000,
+                n_mel_channels=80,
+                mel_norm_file=self.args.mel_norm_file,
+            )
+        else:
+            self.torch_mel_spectrogram_style_encoder = TorchMelSpectrogram(
+                filter_length=4096,
+                hop_length=1024,
+                win_length=4096,
+                normalize=False,
+                sampling_rate=config.audio.sample_rate,
+                mel_fmin=0,
+                mel_fmax=8000,
+                n_mel_channels=80,
+                mel_norm_file=self.args.mel_norm_file,
+            )
+
+        # Load DVAE
+        self.dvae = DiscreteVAE(
+            channels=80,
+            normalization=None,
+            positional_dims=1,
+            num_tokens=self.args.gpt_num_audio_tokens - 2,
+            codebook_dim=512,
+            hidden_dim=512,
+            num_resnet_blocks=3,
+            kernel_size=3,
+            num_layers=2,
+            use_transposed_convs=False,
+        )
+
+        self.dvae.eval()
+        if self.args.dvae_checkpoint:
+            dvae_checkpoint = torch.load(self.args.dvae_checkpoint, map_location=torch.device("cpu"))
+            self.dvae.load_state_dict(dvae_checkpoint, strict=False)
+            print(">> DVAE weights restored from:", self.args.dvae_checkpoint)
+        else:
+            raise RuntimeError(
+                "You need to specify config.model_args.dvae_checkpoint path to be able to train the GPT decoder!!"
+            )
+
+        # Mel spectrogram extractor for DVAE
+        self.torch_mel_spectrogram_dvae = TorchMelSpectrogram(
+            mel_norm_file=self.args.mel_norm_file, sampling_rate=config.audio.dvae_sample_rate
+        )
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def forward(self, text_inputs, text_lengths, audio_codes, wav_lengths, cond_mels, cond_idxs, cond_lens):
+        """
+        Forward pass that uses both text and voice in either text conditioning mode or voice conditioning mode
+        (actuated by `text_first`).
+
+        text_inputs: long tensor, (b,t)
+        text_lengths: long tensor, (b,)
+        mel_inputs:  long tensor, (b,m)
+        wav_lengths: long tensor, (b,)
+        cond_mels: MEL float tensor, (b, num_samples, 80,t_m)
+        cond_idxs: cond start and end indexs, (b, 2)
+        cond_lens: long tensor, (b,)
+        """
+        losses = self.xtts.gpt(
+            text_inputs,
+            text_lengths,
+            audio_codes,
+            wav_lengths,
+            cond_mels=cond_mels,
+            cond_idxs=cond_idxs,
+            cond_lens=cond_lens,
+        )
+        return losses
+
+    @torch.no_grad()
+    def test_run(self, assets) -> Tuple[Dict, Dict]:  # pylint: disable=W0613
+        test_audios = {}
+        if self.config.test_sentences:
+            # init gpt for inference mode
+            self.xtts.gpt.init_gpt_for_inference(kv_cache=self.args.kv_cache, use_deepspeed=False)
+            self.xtts.gpt.eval()
+            print(" | > Synthesizing test sentences.")
+            for idx, s_info in enumerate(self.config.test_sentences):
+                wav = self.xtts.synthesize(
+                    s_info["text"],
+                    self.config,
+                    s_info["speaker_wav"],
+                    s_info["language"],
+                    gpt_cond_len=3,
+                )["wav"]
+                test_audios["{}-audio".format(idx)] = wav
+
+            # delete inference layers
+            del self.xtts.gpt.gpt_inference
+            del self.xtts.gpt.gpt.wte
+        return {"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.args.output_sample_rate)
+
+    def format_batch(self, batch: Dict) -> Dict:
+        return batch
+
+    @torch.no_grad()  # torch no grad to avoid gradients from the pre-processing and DVAE codes extraction
+    def format_batch_on_device(self, batch):
+        """Compute spectrograms on the device."""
+        batch["text_lengths"] = batch["text_lengths"]
+        batch["wav_lengths"] = batch["wav_lengths"]
+        batch["text_inputs"] = batch["padded_text"]
+        batch["cond_idxs"] = batch["cond_idxs"]
+        # compute conditioning mel specs
+        # transform waves from torch.Size([B, num_cond_samples, 1, T] to torch.Size([B * num_cond_samples, 1, T] because if is faster than iterate the tensor
+        B, num_cond_samples, C, T = batch["conditioning"].size()
+        conditioning_reshaped = batch["conditioning"].view(B * num_cond_samples, C, T)
+        paired_conditioning_mel = self.torch_mel_spectrogram_style_encoder(conditioning_reshaped)
+        # transform torch.Size([B * num_cond_samples, n_mel, T_mel]) in torch.Size([B, num_cond_samples, n_mel, T_mel])
+        n_mel = self.torch_mel_spectrogram_style_encoder.n_mel_channels  # paired_conditioning_mel.size(1)
+        T_mel = paired_conditioning_mel.size(2)
+        paired_conditioning_mel = paired_conditioning_mel.view(B, num_cond_samples, n_mel, T_mel)
+        # get the conditioning embeddings
+        batch["cond_mels"] = paired_conditioning_mel
+        # compute codes using DVAE
+        if self.config.audio.sample_rate != self.config.audio.dvae_sample_rate:
+            dvae_wav = torchaudio.functional.resample(
+                batch["wav"],
+                orig_freq=self.config.audio.sample_rate,
+                new_freq=self.config.audio.dvae_sample_rate,
+                lowpass_filter_width=64,
+                rolloff=0.9475937167399596,
+                resampling_method="kaiser_window",
+                beta=14.769656459379492,
+            )
+        else:
+            dvae_wav = batch["wav"]
+        dvae_mel_spec = self.torch_mel_spectrogram_dvae(dvae_wav)
+        codes = self.dvae.get_codebook_indices(dvae_mel_spec)
+
+        batch["audio_codes"] = codes
+        # delete useless batch tensors
+        del batch["padded_text"]
+        del batch["wav"]
+        del batch["conditioning"]
+        return batch
+
+    def train_step(self, batch, criterion):
+        loss_dict = {}
+        cond_mels = batch["cond_mels"]
+        text_inputs = batch["text_inputs"]
+        text_lengths = batch["text_lengths"]
+        audio_codes = batch["audio_codes"]
+        wav_lengths = batch["wav_lengths"]
+        cond_idxs = batch["cond_idxs"]
+        cond_lens = batch["cond_lens"]
+
+        loss_text, loss_mel, _ = self.forward(
+            text_inputs, text_lengths, audio_codes, wav_lengths, cond_mels, cond_idxs, cond_lens
+        )
+        loss_dict["loss_text_ce"] = loss_text * self.args.gpt_loss_text_ce_weight
+        loss_dict["loss_mel_ce"] = loss_mel * self.args.gpt_loss_mel_ce_weight
+        loss_dict["loss"] = loss_dict["loss_text_ce"] + loss_dict["loss_mel_ce"]
+        return {"model_outputs": None}, loss_dict
+
+    def eval_step(self, batch, criterion):
+        # ignore masking for more consistent evaluation
+        batch["cond_idxs"] = None
+        return self.train_step(batch, criterion)
+
+    def on_train_epoch_start(self, trainer):
+        trainer.model.eval()  # the whole model to eval
+        # put gpt model in training mode
+        if hasattr(trainer.model, "module") and hasattr(trainer.model.module, "xtts"):
+            trainer.model.module.xtts.gpt.train()
+        else:
+            trainer.model.xtts.gpt.train()
+
+    def on_init_end(self, trainer):  # pylint: disable=W0613
+        # ignore similarities.pth on clearml save/upload
+        if self.config.dashboard_logger.lower() == "clearml":
+            from clearml.binding.frameworks import WeightsFileHandler
+
+            WeightsFileHandler.add_pre_callback(callback_clearml_load_save)
+
+    @torch.no_grad()
+    def inference(
+        self,
+        x,
+        aux_input=None,
+    ):  # pylint: disable=dangerous-default-value
+        return None
+
+    @staticmethod
+    def get_criterion():
+        return None
+
+    def get_sampler(self, dataset: TTSDataset, num_gpus=1):
+        # sampler for DDP
+        batch_sampler = DistributedSampler(dataset) if num_gpus > 1 else None
+        return batch_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":  # pylint: disable=W0613
+        if is_eval and not config.run_eval:
+            loader = None
+        else:
+            # init dataloader
+            dataset = XTTSDataset(self.config, samples, self.xtts.tokenizer, config.audio.sample_rate, is_eval)
+
+            # wait all the DDP process to be ready
+            if num_gpus > 1:
+                torch.distributed.barrier()
+
+            # sort input sequences from short to long
+            # dataset.preprocess_samples()
+
+            # get samplers
+            sampler = self.get_sampler(dataset, num_gpus)
+
+            # ignore sampler when is eval because if we changed the sampler parameter we will not be able to compare previous runs
+            if sampler is None or is_eval:
+                loader = DataLoader(
+                    dataset,
+                    batch_size=config.eval_batch_size if is_eval else config.batch_size,
+                    shuffle=False,
+                    drop_last=False,
+                    collate_fn=dataset.collate_fn,
+                    num_workers=config.num_eval_loader_workers if is_eval else config.num_loader_workers,
+                    pin_memory=False,
+                )
+            else:
+                loader = DataLoader(
+                    dataset,
+                    sampler=sampler,
+                    batch_size = config.eval_batch_size if is_eval else config.batch_size,
+                    collate_fn=dataset.collate_fn,
+                    num_workers=config.num_eval_loader_workers if is_eval else config.num_loader_workers,
+                    pin_memory=False,
+                )
+        return loader
+
+    def get_optimizer(self) -> List:
+        """Initiate and return the optimizer based on the config parameters."""
+        # ToDo: deal with multi GPU training
+        if self.config.optimizer_wd_only_on_weights:
+            # parameters to only GPT model
+            net = self.xtts.gpt
+
+            # normalizations
+            norm_modules = (
+                nn.BatchNorm2d,
+                nn.InstanceNorm2d,
+                nn.BatchNorm1d,
+                nn.InstanceNorm1d,
+                nn.BatchNorm3d,
+                nn.InstanceNorm3d,
+                nn.GroupNorm,
+                nn.LayerNorm,
+            )
+            # nn.Embedding
+            emb_modules = (nn.Embedding, nn.EmbeddingBag)
+
+            param_names_notweights = set()
+            all_param_names = set()
+            param_map = {}
+            for mn, m in net.named_modules():
+                for k, v in m.named_parameters():
+                    v.is_bias = k.endswith(".bias")
+                    v.is_weight = k.endswith(".weight")
+                    v.is_norm = isinstance(m, norm_modules)
+                    v.is_emb = isinstance(m, emb_modules)
+
+                    fpn = "%s.%s" % (mn, k) if mn else k  # full param name
+                    all_param_names.add(fpn)
+                    param_map[fpn] = v
+                    if v.is_bias or v.is_norm or v.is_emb:
+                        param_names_notweights.add(fpn)
+
+            params_names_notweights = sorted(list(param_names_notweights))
+            params_notweights = [param_map[k] for k in params_names_notweights]
+            params_names_weights = sorted(list(all_param_names ^ param_names_notweights))
+            params_weights = [param_map[k] for k in params_names_weights]
+
+            groups = [
+                {"params": params_weights, "weight_decay": self.config.optimizer_params["weight_decay"]},
+                {"params": params_notweights, "weight_decay": 0},
+            ]
+            # torch.optim.AdamW
+            opt = get_optimizer(
+                self.config.optimizer,
+                self.config.optimizer_params,
+                self.config.lr,
+                parameters=groups,
+            )
+            opt._group_names = [params_names_weights, params_names_notweights]
+            return opt
+
+        return get_optimizer(
+            self.config.optimizer,
+            self.config.optimizer_params,
+            self.config.lr,
+            # optimize only for the GPT model
+            parameters=self.xtts.gpt.parameters(),
+        )
+
+    def get_scheduler(self, optimizer) -> List:
+        """Set the scheduler for the optimizer.
+
+        Args:
+            optimizer: `torch.optim.Optimizer`.
+        """
+        return get_scheduler(self.config.lr_scheduler, self.config.lr_scheduler_params, optimizer)
+
+    def load_checkpoint(
+        self,
+        config,
+        checkpoint_path,
+        eval=False,
+        strict=True,
+        cache_storage="/tmp/tts_cache",
+        target_protocol="s3",
+        target_options={"anon": True},
+    ):  # pylint: disable=unused-argument, disable=W0201, disable=W0102, redefined-builtin
+        """Load the model checkpoint and setup for training or inference"""
+
+        state = self.xtts.get_compatible_checkpoint_state_dict(checkpoint_path)
+
+        # load the model weights
+        self.xtts.load_state_dict(state, strict=strict)
+
+        if eval:
+            self.xtts.gpt.init_gpt_for_inference(kv_cache=self.args.kv_cache, use_deepspeed=False)
+            self.eval()
+            assert not self.training
+
+    @staticmethod
+    def init_from_config(config: "GPTTrainerConfig", samples: Union[List[List], List[Dict]] = None):
+        """Initiate model from config
+
+        Args:
+            config (GPTTrainerConfig): Model config.
+            samples (Union[List[List], List[Dict]]): Training samples to parse speaker ids for training.
+                Defaults to None.
+        """
+        return GPTTrainer(config)

TTS/tts/layers/xtts/xtts_manager.py

@@ -0,0 +1,34 @@
+import torch
+
+class SpeakerManager():
+    def __init__(self, speaker_file_path=None):
+        self.speakers = torch.load(speaker_file_path)
+
+    @property
+    def name_to_id(self):
+        return self.speakers.keys()
+    
+    @property
+    def num_speakers(self):
+        return len(self.name_to_id)
+    
+    @property
+    def speaker_names(self):
+        return list(self.name_to_id.keys())
+    
+
+class LanguageManager():
+    def __init__(self, config):
+        self.langs = config["languages"]
+
+    @property
+    def name_to_id(self):
+        return self.langs
+    
+    @property
+    def num_languages(self):
+        return len(self.name_to_id)
+    
+    @property
+    def language_names(self):
+        return list(self.name_to_id)