update

fireredtts/fireredtts.py

@@ -0,0 +1,163 @@
+import os
+import json
+import torch
+from fireredtts.modules.gpt.gpt import GPT
+from fireredtts.modules import Token2Wav, MelSpectrogramExtractor
+from fireredtts.modules.tokenizer.tokenizer import VoiceBpeTokenizer
+from fireredtts.modules.codec.speaker import SpeakerEmbedddingExtractor
+from fireredtts.utils.utils import load_audio
+
+import time
+
+
+class FireRedTTS:
+    def __init__(self, config_path, pretrained_path, device="cuda"):
+        self.device = device
+        self.config = json.load(open(config_path))
+        self.gpt_path = os.path.join(pretrained_path, "fireredtts_gpt.pt")
+        self.token2wav_path = os.path.join(pretrained_path, "fireredtts_token2wav.pt")
+        self.speaker_extractor_path = os.path.join(
+            pretrained_path, "fireredtts_speaker.bin"
+        )
+        assert os.path.exists(self.token2wav_path)
+        assert os.path.exists(self.gpt_path)
+        assert os.path.exists(self.speaker_extractor_path)
+
+        # tokenizer;
+        self.text_tokenizer = VoiceBpeTokenizer()
+
+        # speaker ectractor
+        self.speaker_extractor = SpeakerEmbedddingExtractor(
+            ckpt_path=self.speaker_extractor_path, device=device
+        )
+
+        # load gpt model
+        self.gpt = GPT(
+            start_text_token=self.config["gpt"]["gpt_start_text_token"],
+            stop_text_token=self.config["gpt"]["gpt_stop_text_token"],
+            layers=self.config["gpt"]["gpt_layers"],
+            model_dim=self.config["gpt"]["gpt_n_model_channels"],
+            heads=self.config["gpt"]["gpt_n_heads"],
+            max_text_tokens=self.config["gpt"]["gpt_max_text_tokens"],
+            max_mel_tokens=self.config["gpt"]["gpt_max_audio_tokens"],
+            max_prompt_tokens=self.config["gpt"]["gpt_max_prompt_tokens"],
+            code_stride_len=self.config["gpt"]["gpt_code_stride_len"],
+            number_text_tokens=self.config["gpt"]["gpt_number_text_tokens"],
+            num_audio_tokens=self.config["gpt"]["gpt_num_audio_tokens"],
+            start_audio_token=self.config["gpt"]["gpt_start_audio_token"],
+            stop_audio_token=self.config["gpt"]["gpt_stop_audio_token"],
+        )
+
+        sd = torch.load(self.gpt_path, map_location=device)["model"]
+        self.gpt.load_state_dict(sd, strict=True)
+        self.gpt = self.gpt.to(device=device)
+        self.gpt.eval()
+        self.gpt.init_gpt_for_inference(kv_cache=True)
+
+        # mel-spectrogram extractor
+        self.mel_extractor = MelSpectrogramExtractor()
+
+        # load token2wav model
+        self.token2wav = Token2Wav.init_from_config(self.config)
+        sd = torch.load(self.token2wav_path, map_location="cpu")
+        self.token2wav.load_state_dict(sd, strict=True)
+        self.token2wav.generator.remove_weight_norm()
+        self.token2wav.eval()
+        self.token2wav = self.token2wav.to(device)
+
+    def extract_spk_embeddings(self, prompt_wav):
+        _, _, audio_resampled = load_audio(audiopath=prompt_wav, sampling_rate=16000)
+        audio_len = torch.tensor(
+            data=[audio_resampled.shape[1]], dtype=torch.long, requires_grad=False
+        )
+
+        # speaker embeddings [1,512]
+        spk_embeddings = self.speaker_extractor(
+            audio_resampled.to(device="cuda")
+        ).unsqueeze(0)
+
+        return spk_embeddings
+
+    def do_gpt_inference(self, spk_gpt, text_tokens):
+        """_summary_
+
+        Args:
+            spk_gpt (_type_): speaker embeddidng in gpt
+            text_tokens (_type_): text tokens
+        """
+        with torch.no_grad():
+            gpt_codes = self.gpt.generate(
+                cond_latents=spk_gpt,
+                text_inputs=text_tokens,
+                input_tokens=None,
+                do_sample=True,
+                top_p=0.85,
+                top_k=30,
+                temperature=0.75,
+                num_return_sequences=9,
+                num_beams=1,
+                length_penalty=1.0,
+                repetition_penalty=2.0,
+                output_attentions=False,
+            )
+
+        seqs = []
+        EOS_TOKEN = self.config["gpt"]["gpt_stop_audio_token"]
+        for seq in gpt_codes:
+            index = (seq == EOS_TOKEN).nonzero(as_tuple=True)[0][0]
+            seq = seq[:index]
+            seqs.append(seq)
+
+        sorted_seqs = sorted(seqs, key=lambda i: len(i), reverse=False)
+        gpt_codes = sorted_seqs[2].unsqueeze(0)  # [1, len]
+        # sorted_len = [len(l) for l in sorted_seqs]
+        # print("---sorted_len:", sorted_len)
+
+        return gpt_codes
+
+    def synthesize(self, prompt_wav, text, lang="auto"):
+        """_summary_
+
+        Args:
+            prompts_wav (_type_): prompts_wav path
+            text (_type_): text
+            lang (_type_): language of text
+        """
+        # Currently only supports Chinese and English
+        assert lang in ["zh", "en", "auto"]
+        assert os.path.exists(prompt_wav)
+
+        # text to tokens
+        text_tokens = self.text_tokenizer.encode(text=text, lang=lang)
+        text_tokens = torch.IntTensor(text_tokens).unsqueeze(0).to(self.device)
+        assert text_tokens.shape[-1] < 400
+
+        # extract speaker embedding
+        spk_embeddings = self.extract_spk_embeddings(prompt_wav=prompt_wav).unsqueeze(0)
+        with torch.no_grad():
+            spk_gpt = self.gpt.reference_embedding(spk_embeddings)
+
+        # gpt inference
+        gpt_start_time = time.time()
+        gpt_codes = self.do_gpt_inference(spk_gpt=spk_gpt, text_tokens=text_tokens)
+        gpt_end_time = time.time()
+        gpt_dur = gpt_end_time - gpt_start_time
+
+        # prompt mel-spectrogram compute
+        prompt_mel = (
+            self.mel_extractor(wav_path=prompt_wav).unsqueeze(0).to(self.device)
+        )
+        # convert token to waveform (b=1, t)
+        voc_start_time = time.time()
+        rec_wavs = self.token2wav.inference(gpt_codes, prompt_mel, n_timesteps=10)
+        voc_end_time = time.time()
+        voc_dur = voc_end_time - voc_start_time
+        all_dur = voc_end_time - gpt_start_time
+
+        # rtf compute
+        # audio_dur = rec_wavs.shape[-1] / 24000
+        # rtf_gpt = gpt_dur / audio_dur
+        # rtf_voc = voc_dur / audio_dur
+        # rtf_all = all_dur / audio_dur
+
+        return rec_wavs

fireredtts/modules/__init__.py

@@ -0,0 +1,42 @@
+import json
+import torch
+import torch.nn as nn
+from fireredtts.modules.bigvgan import get_bigvgan_backend
+from fireredtts.modules.flow import get_flow_frontend, MelSpectrogramExtractor
+
+
+class Token2Wav(nn.Module):
+    def __init__(
+        self,
+        flow: nn.Module,
+        generator: nn.Module,
+    ):
+        super().__init__()
+        self.flow = flow
+        self.generator = generator
+
+    @torch.no_grad()
+    def inference(
+        self, tokens: torch.Tensor, prompt_mel: torch.Tensor, n_timesteps: int = 10
+    ) -> torch.Tensor:
+        token_len = torch.tensor([tokens.shape[1]], dtype=torch.long).to(tokens.device)
+        prompt_mel_len = torch.tensor([prompt_mel.shape[1]], dtype=torch.long).to(
+            prompt_mel.device
+        )
+        # flow
+        mel = self.flow.inference(
+            token=tokens,
+            token_len=token_len,
+            prompt_mel=prompt_mel,
+            prompt_mel_len=prompt_mel_len,
+            n_timesteps=n_timesteps,
+        )
+        # bigvgan
+        audio = self.generator(mel)  # (b=1, 1, t)
+        return audio.squeeze(1)
+
+    @classmethod
+    def init_from_config(cls, config) -> "Token2Wav":
+        flow = get_flow_frontend(config["flow"])
+        bigvgan = get_bigvgan_backend(config["bigvgan"])
+        return cls(flow, bigvgan)

fireredtts/modules/bigvgan/__init__.py

@@ -0,0 +1,6 @@
+from fireredtts.modules.bigvgan.bigvgan import BigVGAN
+
+
+def get_bigvgan_backend(bigvgan_config):
+    generator = BigVGAN(**bigvgan_config)
+    return generator

fireredtts/modules/bigvgan/activations.py

@@ -0,0 +1,126 @@
+# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license.
+
+
+import torch
+from torch import nn, sin, pow
+from torch.nn import Parameter
+
+
+class Snake(nn.Module):
+    """
+    Implementation of a sine-based periodic activation function
+    Shape:
+        - Input: (B, C, T)
+        - Output: (B, C, T), same shape as the input
+    Parameters:
+        - alpha - trainable parameter
+    References:
+        - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+        https://arxiv.org/abs/2006.08195
+    Examples:
+        >>> a1 = snake(256)
+        >>> x = torch.randn(256)
+        >>> x = a1(x)
+    """
+
+    def __init__(
+        self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False
+    ):
+        """
+        Initialization.
+        INPUT:
+            - in_features: shape of the input
+            - alpha: trainable parameter
+            alpha is initialized to 1 by default, higher values = higher-frequency.
+            alpha will be trained along with the rest of your model.
+        """
+        super(Snake, self).__init__()
+        self.in_features = in_features
+
+        # initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:  # log scale alphas initialized to zeros
+            self.alpha = Parameter(torch.zeros(in_features) * alpha)
+        else:  # linear scale alphas initialized to ones
+            self.alpha = Parameter(torch.ones(in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        """
+        Forward pass of the function.
+        Applies the function to the input elementwise.
+        Snake ∶= x + 1/a * sin^2 (xa)
+        """
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # line up with x to [B, C, T]
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+        x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+        return x
+
+
+class SnakeBeta(nn.Module):
+    """
+    A modified Snake function which uses separate parameters for the magnitude of the periodic components
+    Shape:
+        - Input: (B, C, T)
+        - Output: (B, C, T), same shape as the input
+    Parameters:
+        - alpha - trainable parameter that controls frequency
+        - beta - trainable parameter that controls magnitude
+    References:
+        - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+        https://arxiv.org/abs/2006.08195
+    Examples:
+        >>> a1 = snakebeta(256)
+        >>> x = torch.randn(256)
+        >>> x = a1(x)
+    """
+
+    def __init__(
+        self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False
+    ):
+        """
+        Initialization.
+        INPUT:
+            - in_features: shape of the input
+            - alpha - trainable parameter that controls frequency
+            - beta - trainable parameter that controls magnitude
+            alpha is initialized to 1 by default, higher values = higher-frequency.
+            beta is initialized to 1 by default, higher values = higher-magnitude.
+            alpha will be trained along with the rest of your model.
+        """
+        super(SnakeBeta, self).__init__()
+        self.in_features = in_features
+
+        # initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:  # log scale alphas initialized to zeros
+            self.alpha = Parameter(torch.zeros(in_features) * alpha)
+            self.beta = Parameter(torch.zeros(in_features) * alpha)
+        else:  # linear scale alphas initialized to ones
+            self.alpha = Parameter(torch.ones(in_features) * alpha)
+            self.beta = Parameter(torch.ones(in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+        self.beta.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        """
+        Forward pass of the function.
+        Applies the function to the input elementwise.
+        SnakeBeta ∶= x + 1/b * sin^2 (xa)
+        """
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # line up with x to [B, C, T]
+        beta = self.beta.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+            beta = torch.exp(beta)
+        x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+        return x

fireredtts/modules/bigvgan/alias_free_cuda/activation1d.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.
+#   Licensed under the MIT license.
+
+import torch
+import torch.nn as nn
+from token2wav.alias_free_torch.resample import UpSample1d, DownSample1d
+
+# load fused CUDA kernel: this enables importing anti_alias_activation_cuda
+from token2wav.alias_free_cuda import load
+
+load.load()
+
+
+class FusedAntiAliasActivation(torch.autograd.Function):
+    """
+    Assumes filter size 12, replication padding on upsampling, and logscale alpha/beta parameters as inputs
+    """
+
+    @staticmethod
+    def forward(ctx, inputs, ftr, alpha, beta):
+        import anti_alias_activation_cuda
+
+        activation_results = anti_alias_activation_cuda.forward(
+            inputs, ftr, alpha, beta
+        )
+        return activation_results
+
+    @staticmethod
+    def backward(ctx, output_grads):
+        # TODO: implement bwd pass
+        raise NotImplementedError
+        return output_grads, None, None
+
+
+class Activation1d(nn.Module):
+    def __init__(
+        self,
+        activation,
+        up_ratio: int = 2,
+        down_ratio: int = 2,
+        up_kernel_size: int = 12,
+        down_kernel_size: int = 12,
+        fused: bool = True,
+    ):
+        super().__init__()
+        self.up_ratio = up_ratio
+        self.down_ratio = down_ratio
+        self.act = activation
+        self.upsample = UpSample1d(up_ratio, up_kernel_size)
+        self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+        self.fused = fused  # whether to use fused CUDA kernel or not
+
+    def forward(self, x):
+        if not self.fused:
+            x = self.upsample(x)
+            x = self.act(x)
+            x = self.downsample(x)
+            return x
+        else:
+            if self.act.__class__.__name__ == "Snake":
+                beta = self.act.alpha.data  # snake uses same params for alpha and beta
+            else:
+                beta = (
+                    self.act.beta.data
+                )  # snakebeta uses different params for alpha and beta
+            alpha = self.act.alpha.data
+            if (
+                not self.act.alpha_logscale
+            ):  # exp baked into cuda kernel, cancel it out with a log
+                alpha = torch.log(alpha)
+                beta = torch.log(beta)
+            x = FusedAntiAliasActivation.apply(x, self.upsample.filter, alpha, beta)
+            x = self.downsample(x)
+            return x

fireredtts/modules/bigvgan/alias_free_cuda/anti_alias_activation.cpp

@@ -0,0 +1,48 @@
+/* coding=utf-8
+ * Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cuda_fp16.h>
+#include <torch/extension.h>
+#include <vector>
+
+namespace anti_alias_activation {
+
+torch::Tensor fwd_cuda(torch::Tensor const& input,
+                       torch::Tensor const& filter,
+                       torch::Tensor const& alpha,
+		       torch::Tensor const& beta
+		       );
+
+torch::Tensor fwd(torch::Tensor const& input,
+                  torch::Tensor const& filter,
+                  torch::Tensor const& alpha,
+		  torch::Tensor const& beta
+		  ) {
+  AT_ASSERTM(input.dim() == 3, "expected 3D tensor");
+  //AT_ASSERTM((input.scalar_type() == at::ScalarType::Half) ||
+  //	     (input.scalar_type() == at::ScalarType::BFloat16),
+  //    "Only fp16 and bf16 are supported");
+
+  return fwd_cuda(input, filter, alpha, beta);
+}
+
+} // end namespace anti_alias_activation
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward",
+        &anti_alias_activation::fwd,
+	"Anti Alias Activation -- Forward.");
+}

fireredtts/modules/bigvgan/alias_free_cuda/anti_alias_activation_cuda.cu

@@ -0,0 +1,314 @@
+/* coding=utf-8
+ * Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <ATen/ATen.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_profiler_api.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+#include "type_shim.h"
+#include <assert.h>
+#include <cfloat>
+#include <limits>
+#include <stdint.h>
+#include <c10/macros/Macros.h>
+
+namespace {
+
+    /*
+template <typename Datatype, int ELEMENTS_PER_LDG>
+__device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src);
+
+template <>
+__device__ __inline__ void copy_vector<c10::BFloat16, 1>(c10::BFloat16 *dst, const c10::BFloat16 *src) { *dst = *src; }
+
+template <>
+__device__ __inline__ void copy_vector<c10::BFloat16, 4>(c10::BFloat16 *dst, const c10::BFloat16 *src) { *((float2*) dst) = *((float2*) src); }
+
+template <>
+__device__ __inline__ void copy_vector<c10::Half, 1>(c10::Half *dst, const c10::Half *src) { *dst = *src; }
+
+template <>
+__device__ __inline__ void copy_vector<c10::Half, 4>(c10::Half *dst, const c10::Half *src) { *((float2*) dst) = *((float2*) src); }
+
+template <>
+__device__ __inline__ void copy_vector<uint8_t, 1>(uint8_t *dst, const uint8_t *src) { *dst = *src; }
+
+template <>
+__device__ __inline__ void copy_vector<uint8_t, 4>(uint8_t *dst, const uint8_t *src) {*((half2*) dst) = *((half2*) src); }
+
+int log2_ceil(int value) {
+    int log2_value = 0;
+    while ((1 << log2_value) < value) ++log2_value;
+    return log2_value;
+}
+
+template<typename T>
+struct Add {
+  __device__ __forceinline__ T operator()(T a, T b) const {
+    return a + b;
+  }
+};
+
+template<typename T>
+struct Max {
+  __device__ __forceinline__ T operator()(T a, T b) const {
+    return a < b ? b : a;
+  }
+};
+
+template <typename T>
+__device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff)
+{
+#if CUDA_VERSION >= 9000
+    return __shfl_xor_sync(mask, value, laneMask, width);
+#else
+    return __shfl_xor(value, laneMask, width);
+#endif
+}
+
+template <typename acc_t, int WARP_BATCH, int WARP_SIZE, template<typename> class ReduceOp>
+__device__ __forceinline__ void warp_reduce(acc_t* sum) {
+    ReduceOp<acc_t> r;
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        #pragma unroll
+        for (int i = 0;  i < WARP_BATCH;  ++i) {
+            acc_t b = WARP_SHFL_XOR_NATIVE(sum[i], offset, WARP_SIZE);
+            sum[i] = r(sum[i], b);
+        }
+    }
+}
+*/
+
+template <typename input_t, typename output_t, typename acc_t>
+__global__ void anti_alias_activation_forward(
+    output_t *dst,
+    const input_t *src,
+    const input_t *ftr,
+    const input_t *alpha,
+    const input_t *beta,
+    int batch_size,
+    int channels,
+    int seq_len)
+{
+    // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and
+    constexpr int ELEMENTS_PER_LDG_STG = 1; //(WARP_ITERATIONS < 4) ? 1 : 4;
+    constexpr int BUFFER_SIZE = 32;
+    constexpr int FILTER_SIZE = 12;
+    constexpr int HALF_FILTER_SIZE = 6;
+    constexpr int REPLICATION_PAD = 5; // 5 on each side
+
+    // blockDim/threadIdx = (128, 1, 1)
+    // gridDim/blockIdx = (seq_blocks, channels, batches)
+    int block_offset = (blockIdx.x * 128 * BUFFER_SIZE + seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
+    int local_offset = threadIdx.x * BUFFER_SIZE;
+    int seq_offset = blockIdx.x * 128 * BUFFER_SIZE + local_offset;
+
+
+    //int intermediate_seq_len = seq_len * 2 - 1 + 4 * REPLICATION_PAD;
+    //int intermediate_block_offset = (blockIdx.x * 128 * BUFFER_SIZE * 2 + intermediate_seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
+    //int intermediate_local_offset = threadIdx.x * BUFFER_SIZE * 2;
+
+    int output_seq_len = seq_len * 2 ; //
+    int output_block_offset = (blockIdx.x * 128 * BUFFER_SIZE * 2 + output_seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
+    int output_local_offset = threadIdx.x * BUFFER_SIZE * 2;
+    int output_seq_offset = blockIdx.x * 128 * BUFFER_SIZE *2 + output_local_offset;
+    // get values needed for replication padding before moving pointer
+    const input_t *right_most_pntr = src + (seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
+    input_t seq_left_most_value = right_most_pntr[0];
+    input_t seq_right_most_value = right_most_pntr[seq_len - 1];
+
+    src += block_offset + local_offset;
+    dst += output_block_offset + output_local_offset  ;
+    alpha = alpha + blockIdx.y;
+    input_t alpha_val = expf(alpha[0]);
+    beta = beta + blockIdx.y;
+    input_t beta_val = expf(beta[0]);
+    // load data from global memory
+    input_t elements[2*FILTER_SIZE+2*BUFFER_SIZE] = {0};
+    input_t intermediates[2*FILTER_SIZE+2*BUFFER_SIZE] = {0};
+    //output_t output[2*BUFFER_SIZE];
+    input_t filter[FILTER_SIZE];
+    //input_t temp_data[ELEMENTS_PER_LDG_STG];
+    //uint8_t temp_mask[ELEMENTS_PER_LDG_STG];
+
+    #pragma unroll
+    for (int it = 0; it < FILTER_SIZE; it+=1) {
+        filter[it] = ftr[it];
+    }
+
+
+    #pragma unroll
+    for (int it = -HALF_FILTER_SIZE;  it < BUFFER_SIZE + HALF_FILTER_SIZE ;  it+=1) {
+        int element_index = seq_offset + it;
+	if ((element_index < 0) && (element_index >= -REPLICATION_PAD)) {
+	    elements[2*(HALF_FILTER_SIZE+it)] = 2*seq_left_most_value;
+	}
+	if ((element_index >= seq_len) && (element_index < seq_len + REPLICATION_PAD)) {
+	    elements[2*(HALF_FILTER_SIZE+it)] = 2*seq_right_most_value;
+	}
+        if ((element_index >= 0) && (element_index < seq_len)) {
+	  elements[2*(HALF_FILTER_SIZE+it)] = 2*src[it];
+        }
+    }
+
+
+
+    // apply filter
+    #pragma unroll
+    for (int it = 0;  it < (2 * BUFFER_SIZE + 2*FILTER_SIZE);  it+=1) {
+        input_t acc = 0.0;
+
+	int element_index = output_seq_offset + it; // index for output
+	#pragma unroll
+        for (int f_idx = 0; f_idx < FILTER_SIZE; f_idx+=1){
+	  if ((element_index + f_idx) >= 0){
+            acc += filter[f_idx] * elements[it+f_idx];
+	  }
+	}
+        intermediates[it] = acc;
+    }
+
+    double no_div_by_zero = 0.000000001;
+    #pragma unroll
+    for (int it = 0; it < 12 + 2 * BUFFER_SIZE; it++) {
+        intermediates[it] += (1.0/(beta_val + no_div_by_zero)) *  sinf(intermediates[it] * alpha_val) * sinf(intermediates[it] * alpha_val);
+    }
+
+
+    // now copy to output
+    #pragma unroll
+    for (int it = 0; it < 2*BUFFER_SIZE; it+=1){
+      int element_index = output_seq_offset + it;
+      if (element_index < output_seq_len) {
+	dst[it]  = intermediates[it+6];
+      }
+    }
+
+
+
+    // for (int it = 0;  it < BUFFER_SIZE;  it+=ELEMENTS_PER_LDG_STG) {
+    //     int element_index = seq_offset + it;
+    //     if (element_index < seq_len) {
+    //         dst[it] = output[it];
+    //     }
+    // }
+
+
+    // // Upsample convolution
+    // for (int it = 0;  it < 2 * BUFFER_SIZE + 12;  it+=1) {
+    //     input_t acc = 0.0;
+
+    //     for (int f_idx = 0; f_idx < FILTER_SIZE; f_idx+=1){
+    //         acc += filter[f_idx] * elements[it+f_idx];
+    //     }
+    //     intermediates[it] = acc;
+    // }
+
+    // // correct the corners of intermediates
+    // if (seq_offset == 0) {
+    //     for (int it = 0; it < 6; it+=1)
+    //         intermediates[it] = 0;
+    // }
+
+    // if (seq_offset + 32 >= seq_len) {
+    //     int offset = seq_len % 32 == 0 ? 32 : seq_len % 32;
+
+    //     for (int it = 0; it < 6; it++) {
+    //         intermediates[6+2*offset+it] = 0;
+    //     }
+    // }
+
+
+
+
+    // for (int it = 0;  it < BUFFER_SIZE;  it+=ELEMENTS_PER_LDG_STG) {
+    //     int element_index = seq_offset + it;
+    //     if (element_index < seq_len) {
+    //         dst[it] = output[it];
+    //     }
+    // }
+}
+
+template<typename input_t, typename output_t, typename acc_t>
+void dispatch_anti_alias_activation_forward(
+    output_t *dst,
+    const input_t *src,
+    const input_t *ftr,
+    const input_t *alpha,
+    const input_t *beta,
+    int batch_size,
+    int channels,
+    int seq_len)
+{
+    if (seq_len == 0) {
+        return;
+    } else {
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+        constexpr int seq_len_per_block = 4096;
+        int blocks_per_seq_len = (seq_len + seq_len_per_block - 1) / seq_len_per_block;
+        dim3 blocks(blocks_per_seq_len, channels, batch_size);
+        dim3 threads(threads_per_block, 1, 1);
+
+        anti_alias_activation_forward<input_t, output_t, acc_t>
+	  <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, ftr, alpha, beta, batch_size, channels, seq_len);
+    }
+}
+}
+
+namespace anti_alias_activation {
+
+  torch::Tensor fwd_cuda(torch::Tensor const& input, torch::Tensor const& filter, torch::Tensor const& alpha, torch::Tensor const& beta)
+{
+  // input is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len]
+  const int batches = input.size(0);
+  const int channels = input.size(1);
+  const int seq_len = input.size(2);
+
+  // Output
+  auto act_options = input.options().requires_grad(false);
+  int output_seq_len = seq_len*2; // we'll be dilating between each element by interspersing with zeros
+
+  torch::Tensor anti_alias_activation_results =
+      torch::empty({batches, channels, output_seq_len}, act_options);
+
+  // Softmax Intermediate Result Ptr
+  void* input_ptr = static_cast<void*>(input.data_ptr());
+  void* filter_ptr = static_cast<void*>(filter.data_ptr());
+  void* alpha_ptr = static_cast<void*>(alpha.data_ptr());
+  void* beta_ptr = static_cast<void*>(beta.data_ptr());
+  void* anti_alias_activation_results_ptr = static_cast<void*>(anti_alias_activation_results.data_ptr());
+
+  DISPATCH_FLOAT_HALF_AND_BFLOAT(
+      input.scalar_type(),
+      "dispatch anti alias activation_forward",
+      dispatch_anti_alias_activation_forward<scalar_t, scalar_t, float>(
+        reinterpret_cast<scalar_t*>(anti_alias_activation_results_ptr),
+	    reinterpret_cast<const scalar_t*>(input_ptr),
+        reinterpret_cast<const scalar_t*>(filter_ptr),
+        reinterpret_cast<const scalar_t*>(alpha_ptr),
+	reinterpret_cast<const scalar_t*>(beta_ptr),
+	    batches,
+        channels,
+        seq_len);
+      );
+  return anti_alias_activation_results;
+}
+}

fireredtts/modules/bigvgan/alias_free_cuda/compat.h

@@ -0,0 +1,31 @@
+/* coding=utf-8
+ * Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*This code is copied fron NVIDIA apex:
+ *     https://github.com/NVIDIA/apex
+ *     with minor changes. */
+
+
+
+#ifndef TORCH_CHECK
+#define TORCH_CHECK AT_CHECK
+#endif
+
+#ifdef VERSION_GE_1_3
+#define DATA_PTR data_ptr
+#else
+#define DATA_PTR data
+#endif

fireredtts/modules/bigvgan/alias_free_cuda/load.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.
+#   Licensed under the MIT license.
+
+import os
+import pathlib
+import subprocess
+
+from torch.utils import cpp_extension
+
+# Setting this param to a list has a problem of generating different
+# compilation commands (with diferent order of architectures) and
+# leading to recompilation of fused kernels. Set it to empty string
+# to avoid recompilation and assign arch flags explicity in
+# extra_cuda_cflags below
+os.environ["TORCH_CUDA_ARCH_LIST"] = ""
+
+
+def load():
+    # Check if cuda 11 is installed for compute capability 8.0
+    cc_flag = []
+    _, bare_metal_major, _ = _get_cuda_bare_metal_version(cpp_extension.CUDA_HOME)
+    if int(bare_metal_major) >= 11:
+        cc_flag.append("-gencode")
+        cc_flag.append("arch=compute_80,code=sm_80")
+
+    # Build path
+    srcpath = pathlib.Path(__file__).parent.absolute()
+    buildpath = srcpath / "build"
+    _create_build_dir(buildpath)
+
+    # Helper function to build the kernels.
+    def _cpp_extention_load_helper(name, sources, extra_cuda_flags):
+        return cpp_extension.load(
+            name=name,
+            sources=sources,
+            build_directory=buildpath,
+            extra_cflags=[
+                "-O3",
+            ],
+            extra_cuda_cflags=[
+                "-O3",
+                "-gencode",
+                "arch=compute_70,code=sm_70",
+                "--use_fast_math",
+            ]
+            + extra_cuda_flags
+            + cc_flag,
+            verbose=True,
+        )
+
+    extra_cuda_flags = [
+        "-U__CUDA_NO_HALF_OPERATORS__",
+        "-U__CUDA_NO_HALF_CONVERSIONS__",
+        "--expt-relaxed-constexpr",
+        "--expt-extended-lambda",
+    ]
+
+    sources = [
+        srcpath / "anti_alias_activation.cpp",
+        srcpath / "anti_alias_activation_cuda.cu",
+    ]
+    anti_alias_activation_cuda = _cpp_extention_load_helper(
+        "anti_alias_activation_cuda", sources, extra_cuda_flags
+    )
+
+
+def _get_cuda_bare_metal_version(cuda_dir):
+    raw_output = subprocess.check_output(
+        [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
+    )
+    output = raw_output.split()
+    release_idx = output.index("release") + 1
+    release = output[release_idx].split(".")
+    bare_metal_major = release[0]
+    bare_metal_minor = release[1][0]
+
+    return raw_output, bare_metal_major, bare_metal_minor
+
+
+def _create_build_dir(buildpath):
+    try:
+        os.mkdir(buildpath)
+    except OSError:
+        if not os.path.isdir(buildpath):
+            print(f"Creation of the build directory {buildpath} failed")

fireredtts/modules/bigvgan/alias_free_cuda/test_activation.py

@@ -0,0 +1,64 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.
+#   Licensed under the MIT license.
+
+import math
+import torch
+import alias_free_cuda
+from alias_free_cuda import activation1d
+from activations import Snake, SnakeBeta
+
+
+def test_load_fused_kernels():
+    try:
+        import alias_free_cuda
+        import torch
+
+        print("[Success] load_fused_kernels")
+    except ImportError as e:
+        print("[Fail] load_fused_kernels")
+        raise e
+
+
+def test_anti_alias_activation():
+    data = torch.rand((10, 10, 50000), device="cuda")
+
+    # check activations.Snake cuda vs. torch
+    fused_anti_alias_activation = activation1d.Activation1d(
+        activation=Snake(10), fused=True
+    ).cuda()
+    fused_activation_output = fused_anti_alias_activation(data)
+
+    torch_anti_alias_activation = activation1d.Activation1d(
+        activation=Snake(10), fused=False
+    ).cuda()
+    torch_activation_output = torch_anti_alias_activation(data)
+
+    test_result = (fused_activation_output - torch_activation_output).abs()
+
+    while test_result.dim() != 1:
+        test_result = test_result.mean(dim=-1)
+
+    diff = test_result.mean(dim=-1)
+
+    if diff <= 1e-3:
+        print(
+            f"\n[Success] test_fused_anti_alias_activation"
+            f"\n > mean_difference={diff}"
+            f"\n > fused_values={fused_activation_output[-1][-1][-100:].tolist()}"
+            f"\n > torch_values={torch_activation_output[-1][-1][-100:].tolist()}"
+        )
+    else:
+        print(
+            f"\n[Fail] test_fused_anti_alias_activation"
+            f"\n > mean_difference={diff}, "
+            f"\n > fused_values={fused_activation_output[-1][-1][-30:].tolist()}, "
+            f"\n > torch_values={torch_activation_output[-1][-1][-30:].tolist()}"
+        )
+
+
+if __name__ == "__main__":
+    from alias_free_cuda import load
+
+    load.load()
+    test_load_fused_kernels()
+    test_anti_alias_activation()

fireredtts/modules/bigvgan/alias_free_cuda/test_activation_snake_beta.py

@@ -0,0 +1,64 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.
+#   Licensed under the MIT license.
+
+import math
+import torch
+import alias_free_cuda
+from alias_free_cuda import activation1d
+from activations import Snake, SnakeBeta
+
+
+def test_load_fused_kernels():
+    try:
+        import alias_free_cuda
+        import torch
+
+        print("[Success] load_fused_kernels")
+    except ImportError as e:
+        print("[Fail] load_fused_kernels")
+        raise e
+
+
+def test_anti_alias_activation():
+    data = torch.rand((10, 10, 50000), device="cuda")
+
+    # check activations.Snake cuda vs. torch
+    fused_anti_alias_activation = activation1d.Activation1d(
+        activation=SnakeBeta(10), fused=True
+    ).cuda()
+    fused_activation_output = fused_anti_alias_activation(data)
+
+    torch_anti_alias_activation = activation1d.Activation1d(
+        activation=SnakeBeta(10), fused=False
+    ).cuda()
+    torch_activation_output = torch_anti_alias_activation(data)
+
+    test_result = (fused_activation_output - torch_activation_output).abs()
+
+    while test_result.dim() != 1:
+        test_result = test_result.mean(dim=-1)
+
+    diff = test_result.mean(dim=-1)
+
+    if diff <= 1e-3:
+        print(
+            f"\n[Success] test_fused_anti_alias_activation"
+            f"\n > mean_difference={diff}"
+            f"\n > fused_values={fused_activation_output[-1][-1][-100:].tolist()}"
+            f"\n > torch_values={torch_activation_output[-1][-1][-100:].tolist()}"
+        )
+    else:
+        print(
+            f"\n[Fail] test_fused_anti_alias_activation"
+            f"\n > mean_difference={diff}, "
+            f"\n > fused_values={fused_activation_output[-1][-1][-30:].tolist()}, "
+            f"\n > torch_values={torch_activation_output[-1][-1][-30:].tolist()}"
+        )
+
+
+if __name__ == "__main__":
+    from alias_free_cuda import load
+
+    load.load()
+    test_load_fused_kernels()
+    test_anti_alias_activation()

fireredtts/modules/bigvgan/alias_free_cuda/type_shim.h

@@ -0,0 +1,97 @@
+/* coding=utf-8
+ * Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+
+#include <ATen/ATen.h>
+#include "compat.h"
+
+
+#define DISPATCH_FLOAT_HALF_AND_BFLOAT(TYPE, NAME, ...)			\
+  switch(TYPE)							\
+    {									\
+    case at::ScalarType::Float:			\
+      {									\
+	using scalar_t = float;			    \
+	__VA_ARGS__;						\
+	break;								\
+      }                                 \
+    case at::ScalarType::Half:			\
+      {									\
+	using scalar_t = at::Half;			\
+	__VA_ARGS__;						\
+	break;								\
+      }									\
+    case at::ScalarType::BFloat16:		\
+      {									\
+	using scalar_t = at::BFloat16;		\
+	__VA_ARGS__;						\
+	break;								\
+      }									\
+    default:							\
+      AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'");	\
+      }
+
+
+
+#define DISPATCH_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(TYPEIN, TYPEOUT, NAME, ...) \
+  switch(TYPEIN)							\
+    {									\
+    case at::ScalarType::Float:						\
+      {									\
+	using scalar_t_in = float;					\
+	switch(TYPEOUT)							\
+	  {								\
+	  case at::ScalarType::Float:					\
+	    {								\
+	      using scalar_t_out = float;				\
+	      __VA_ARGS__;						\
+	      break;							\
+	    }								\
+	  case at::ScalarType::Half:					\
+	    {								\
+	      using scalar_t_out = at::Half;				\
+	      __VA_ARGS__;						\
+	      break;							\
+	    }								\
+	  case at::ScalarType::BFloat16:				\
+	    {								\
+	      using scalar_t_out = at::BFloat16;			\
+	      __VA_ARGS__;						\
+	      break;							\
+	    }								\
+	  default:							\
+	    AT_ERROR(#NAME, " not implemented for '", toString(TYPEOUT), "'"); \
+	  }								\
+	break;								\
+      }									\
+    case at::ScalarType::Half:						\
+      {									\
+	using scalar_t_in = at::Half;					\
+	using scalar_t_out = at::Half;					\
+	__VA_ARGS__;							\
+	break;								\
+      }									\
+    case at::ScalarType::BFloat16:					\
+      {									\
+	using scalar_t_in = at::BFloat16;				\
+	using scalar_t_out = at::BFloat16;				\
+	__VA_ARGS__;							\
+	break;								\
+      }									\
+    default:								\
+      AT_ERROR(#NAME, " not implemented for '", toString(TYPEIN), "'");	\
+    }
+

fireredtts/modules/bigvgan/alias_free_torch/__init__.py

@@ -0,0 +1,5 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+
+from .filter import *
+from .resample import *
+from .act import *

fireredtts/modules/bigvgan/alias_free_torch/act.py

@@ -0,0 +1,29 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+
+import torch.nn as nn
+from .resample import UpSample1d, DownSample1d
+
+
+class Activation1d(nn.Module):
+    def __init__(
+        self,
+        activation,
+        up_ratio: int = 2,
+        down_ratio: int = 2,
+        up_kernel_size: int = 12,
+        down_kernel_size: int = 12,
+    ):
+        super().__init__()
+        self.up_ratio = up_ratio
+        self.down_ratio = down_ratio
+        self.act = activation
+        self.upsample = UpSample1d(up_ratio, up_kernel_size)
+        self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+    # x: [B,C,T]
+    def forward(self, x):
+        x = self.upsample(x)
+        x = self.act(x)
+        x = self.downsample(x)
+
+        return x

fireredtts/modules/bigvgan/alias_free_torch/filter.py

@@ -0,0 +1,98 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+if "sinc" in dir(torch):
+    sinc = torch.sinc
+else:
+    # This code is adopted from adefossez's julius.core.sinc under the MIT License
+    # https://adefossez.github.io/julius/julius/core.html
+    #   LICENSE is in incl_licenses directory.
+    def sinc(x: torch.Tensor):
+        """
+        Implementation of sinc, i.e. sin(pi * x) / (pi * x)
+        __Warning__: Different to julius.sinc, the input is multiplied by `pi`!
+        """
+        return torch.where(
+            x == 0,
+            torch.tensor(1.0, device=x.device, dtype=x.dtype),
+            torch.sin(math.pi * x) / math.pi / x,
+        )
+
+
+# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License
+# https://adefossez.github.io/julius/julius/lowpass.html
+#   LICENSE is in incl_licenses directory.
+def kaiser_sinc_filter1d(
+    cutoff, half_width, kernel_size
+):  # return filter [1,1,kernel_size]
+    even = kernel_size % 2 == 0
+    half_size = kernel_size // 2
+
+    # For kaiser window
+    delta_f = 4 * half_width
+    A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+    if A > 50.0:
+        beta = 0.1102 * (A - 8.7)
+    elif A >= 21.0:
+        beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0)
+    else:
+        beta = 0.0
+    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+
+    # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio
+    if even:
+        time = torch.arange(-half_size, half_size) + 0.5
+    else:
+        time = torch.arange(kernel_size) - half_size
+    if cutoff == 0:
+        filter_ = torch.zeros_like(time)
+    else:
+        filter_ = 2 * cutoff * window * sinc(2 * cutoff * time)
+        # Normalize filter to have sum = 1, otherwise we will have a small leakage
+        # of the constant component in the input signal.
+        filter_ /= filter_.sum()
+        filter = filter_.view(1, 1, kernel_size)
+
+    return filter
+
+
+class LowPassFilter1d(nn.Module):
+    def __init__(
+        self,
+        cutoff=0.5,
+        half_width=0.6,
+        stride: int = 1,
+        padding: bool = True,
+        padding_mode: str = "replicate",
+        kernel_size: int = 12,
+    ):
+        # kernel_size should be even number for stylegan3 setup,
+        # in this implementation, odd number is also possible.
+        super().__init__()
+        if cutoff < -0.0:
+            raise ValueError("Minimum cutoff must be larger than zero.")
+        if cutoff > 0.5:
+            raise ValueError("A cutoff above 0.5 does not make sense.")
+        self.kernel_size = kernel_size
+        self.even = kernel_size % 2 == 0
+        self.pad_left = kernel_size // 2 - int(self.even)
+        self.pad_right = kernel_size // 2
+        self.stride = stride
+        self.padding = padding
+        self.padding_mode = padding_mode
+        filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
+        self.register_buffer("filter", filter)
+
+    # input [B, C, T]
+    def forward(self, x):
+        _, C, _ = x.shape
+
+        if self.padding:
+            x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+        out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
+
+        return out

fireredtts/modules/bigvgan/alias_free_torch/resample.py

@@ -0,0 +1,57 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+
+import torch.nn as nn
+from torch.nn import functional as F
+from .filter import LowPassFilter1d
+from .filter import kaiser_sinc_filter1d
+
+
+class UpSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = (
+            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        )
+        self.stride = ratio
+        self.pad = self.kernel_size // ratio - 1
+        self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+        self.pad_right = (
+            self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+        )
+        filter = kaiser_sinc_filter1d(
+            cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size
+        )
+        self.register_buffer("filter", filter)
+
+    # x: [B, C, T]
+    def forward(self, x):
+        _, C, _ = x.shape
+
+        x = F.pad(x, (self.pad, self.pad), mode="replicate")
+        x = self.ratio * F.conv_transpose1d(
+            x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C
+        )
+        x = x[..., self.pad_left : -self.pad_right]
+
+        return x
+
+
+class DownSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = (
+            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        )
+        self.lowpass = LowPassFilter1d(
+            cutoff=0.5 / ratio,
+            half_width=0.6 / ratio,
+            stride=ratio,
+            kernel_size=self.kernel_size,
+        )
+
+    def forward(self, x):
+        xx = self.lowpass(x)
+
+        return xx

fireredtts/modules/bigvgan/bigvgan.py

@@ -0,0 +1,399 @@
+import typing as tp
+import torch
+import torch.nn as nn
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+from fireredtts.modules.bigvgan.alias_free_torch import (
+    Activation1d as TorchActivation1d,
+)
+from fireredtts.modules.bigvgan.activations import Snake, SnakeBeta
+
+
+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 AMPBlock1(torch.nn.Module):
+    def __init__(
+        self,
+        channels,
+        kernel_size=3,
+        dilation=(1, 3, 5),
+        activation=None,
+        snake_logscale=True,
+        use_cuda_kernel=False,
+    ):
+        super(AMPBlock1, self).__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.convs1.apply(init_weights)
+
+        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),
+                    )
+                ),
+            ]
+        )
+        self.convs2.apply(init_weights)
+
+        self.num_layers = len(self.convs1) + len(
+            self.convs2
+        )  # total number of conv layers
+
+        # select which Activation1d, lazy-load cuda version to ensure backward compatibility
+        if use_cuda_kernel:
+            from modules.bigvgan.alias_free_cuda.activation1d import (
+                Activation1d as CudaActivation1d,
+            )
+
+            Activation1d = CudaActivation1d
+        else:
+            Activation1d = TorchActivation1d
+
+        if (
+            activation == "snake"
+        ):  # periodic nonlinearity with snake function and anti-aliasing
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=Snake(channels, alpha_logscale=snake_logscale)
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+        elif (
+            activation == "snakebeta"
+        ):  # periodic nonlinearity with snakebeta function and anti-aliasing
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=SnakeBeta(channels, alpha_logscale=snake_logscale)
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+        else:
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
+
+    def forward(self, x):
+        acts1, acts2 = self.activations[::2], self.activations[1::2]
+        for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2):
+            xt = a1(x)
+            xt = c1(xt)
+            xt = a2(xt)
+            xt = c2(xt)
+            x = xt + x
+
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class AMPBlock2(torch.nn.Module):
+    def __init__(
+        self,
+        channels,
+        kernel_size=3,
+        dilation=(1, 3),
+        activation=None,
+        snake_logscale=True,
+        use_cuda_kernel=False,
+    ):
+        super(AMPBlock2, self).__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]),
+                    )
+                ),
+            ]
+        )
+        self.convs.apply(init_weights)
+
+        self.num_layers = len(self.convs)  # total number of conv layers
+
+        # select which Activation1d, lazy-load cuda version to ensure backward compatibility
+        if use_cuda_kernel:
+            from modules.bigvgan.alias_free_cuda.activation1d import (
+                Activation1d as CudaActivation1d,
+            )
+
+            Activation1d = CudaActivation1d
+        else:
+            Activation1d = TorchActivation1d
+
+        if (
+            activation == "snake"
+        ):  # periodic nonlinearity with snake function and anti-aliasing
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=Snake(channels, alpha_logscale=snake_logscale)
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+        elif (
+            activation == "snakebeta"
+        ):  # periodic nonlinearity with snakebeta function and anti-aliasing
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=SnakeBeta(channels, alpha_logscale=snake_logscale)
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+        else:
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
+
+    def forward(self, x):
+        for c, a in zip(self.convs, self.activations):
+            xt = a(x)
+            xt = c(xt)
+            x = xt + x
+
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+class BigVGAN(torch.nn.Module):
+    # this is our main BigVGAN model. Applies anti-aliased periodic activation for resblocks.
+    def __init__(
+        self,
+        num_mels: int,
+        upsample_initial_channel: int,
+        resblock_kernel_sizes: tp.List[int],
+        resblock_dilation_sizes: tp.List[tp.List[int]],
+        upsample_rates: tp.List[int],
+        upsample_kernel_sizes: tp.List[int],
+        resblock_type: str = "1",
+        snake_logscale: bool = True,
+        activation: str = "snakebeta",
+        use_tanh_at_final: bool = False,
+        use_bias_at_final: bool = False,
+        use_cuda_kernel: bool = False,
+    ):
+        super(BigVGAN, self).__init__()
+
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+
+        # pre conv
+        self.conv_pre = weight_norm(
+            Conv1d(num_mels, upsample_initial_channel, 7, 1, padding=3)
+        )
+
+        # define which AMPBlock to use. BigVGAN uses AMPBlock1 as default
+        resblock = AMPBlock1 if resblock_type == "1" else AMPBlock2
+
+        # transposed conv-based upsamplers. does not apply anti-aliasing
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(
+                nn.ModuleList(
+                    [
+                        weight_norm(
+                            ConvTranspose1d(
+                                upsample_initial_channel // (2**i),
+                                upsample_initial_channel // (2 ** (i + 1)),
+                                k,
+                                u,
+                                padding=(k - u) // 2,
+                            )
+                        )
+                    ]
+                )
+            )
+
+        # residual blocks using anti-aliased multi-periodicity composition modules (AMP)
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(
+                zip(resblock_kernel_sizes, resblock_dilation_sizes)
+            ):
+                self.resblocks.append(
+                    resblock(
+                        ch,
+                        k,
+                        d,
+                        activation=activation,
+                        snake_logscale=snake_logscale,
+                        use_cuda_kernel=use_cuda_kernel,
+                    )
+                )
+
+        # select which Activation1d, lazy-load cuda version to ensure backward compatibility
+        if use_cuda_kernel:
+            from modules.bigvgan.alias_free_cuda.activation1d import (
+                Activation1d as CudaActivation1d,
+            )
+
+            Activation1d = CudaActivation1d
+        else:
+            Activation1d = TorchActivation1d
+
+        # post conv
+        if (
+            activation == "snake"
+        ):  # periodic nonlinearity with snake function and anti-aliasing
+            activation_post = Snake(ch, alpha_logscale=snake_logscale)
+            self.activation_post = Activation1d(activation=activation_post)
+        elif (
+            activation == "snakebeta"
+        ):  # periodic nonlinearity with snakebeta function and anti-aliasing
+            activation_post = SnakeBeta(ch, alpha_logscale=snake_logscale)
+            self.activation_post = Activation1d(activation=activation_post)
+        else:
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
+
+        # whether to use bias for the final conv_post. Defaults to True for backward compatibility
+        self.use_bias_at_final = use_bias_at_final
+        self.conv_post = weight_norm(
+            Conv1d(ch, 1, 7, 1, padding=3, bias=self.use_bias_at_final)
+        )
+
+        # weight initialization
+        for i in range(len(self.ups)):
+            self.ups[i].apply(init_weights)
+        self.conv_post.apply(init_weights)
+
+        # final tanh activation. Defaults to True for backward compatibility
+        self.use_tanh_at_final = use_tanh_at_final
+
+    def forward(self, x):
+        # pre conv
+        x = self.conv_pre(x)
+
+        for i in range(self.num_upsamples):
+            # upsampling
+            for i_up in range(len(self.ups[i])):
+                x = self.ups[i][i_up](x)
+            # AMP blocks
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+
+        # post conv
+        x = self.activation_post(x)
+        x = self.conv_post(x)
+        # final tanh activation
+        if self.use_tanh_at_final:
+            x = torch.tanh(x)
+        else:
+            x = torch.clamp(x, min=-1.0, max=1.0)  # bound the output to [-1, 1]
+
+        return x
+
+    def remove_weight_norm(self):
+        print("Removing weight norm...")
+        for l in self.ups:
+            for l_i in l:
+                remove_weight_norm(l_i)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+        remove_weight_norm(self.conv_pre)
+        remove_weight_norm(self.conv_post)

fireredtts/modules/codec/speaker.py

@@ -0,0 +1,1052 @@
+from librosa.filters import mel as librosa_mel_fn
+from torch import nn
+from torch.nn import functional as F
+
+import math
+import numpy as np
+import torch
+import torchaudio
+
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def spectral_normalize_torch(magnitudes):
+    output = dynamic_range_compression_torch(magnitudes)
+    return output
+
+
+class TorchMelSpectrogram(nn.Module):
+
+    def __init__(
+        self,
+        filter_length=1024,
+        hop_length=160,
+        win_length=640,
+        n_mel_channels=80,
+        mel_fmin=0,
+        mel_fmax=8000,
+        sampling_rate=16000,
+    ):
+
+        super().__init__()
+        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_basis = {}
+        self.hann_window = {}
+
+    def forward(self, inp, length=None):
+        if len(inp.shape) == 3:
+            inp = inp.squeeze(1) if inp.shape[1] == 1 else inp.squeeze(2)
+        assert len(inp.shape) == 2
+
+        y = inp
+        if len(list(self.mel_basis.keys())) == 0:
+            mel = librosa_mel_fn(
+                sr=self.sampling_rate,
+                n_fft=self.filter_length,
+                n_mels=self.n_mel_channels,
+                fmin=self.mel_fmin,
+                fmax=self.mel_fmax,
+            )
+            self.mel_basis[str(self.mel_fmax) + "_" + str(y.device)] = (
+                torch.from_numpy(mel).float().to(y.device)
+            )
+            self.hann_window[str(y.device)] = torch.hann_window(self.win_length).to(
+                y.device
+            )
+
+        y = torch.nn.functional.pad(
+            y.unsqueeze(1),
+            (
+                int((self.filter_length - self.hop_length) / 2),
+                int((self.filter_length - self.hop_length) / 2),
+            ),
+            mode="reflect",
+        )
+        y = y.squeeze(1)
+
+        # complex tensor as default, then use view_as_real for future pytorch compatibility
+        spec = torch.stft(
+            y,
+            self.filter_length,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            window=self.hann_window[str(y.device)],
+            center=False,
+            pad_mode="reflect",
+            normalized=False,
+            onesided=True,
+            return_complex=True,
+        )
+        spec = torch.view_as_real(spec)
+        spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9))
+
+        spec = torch.matmul(
+            self.mel_basis[str(self.mel_fmax) + "_" + str(y.device)], spec
+        )
+        spec = spectral_normalize_torch(spec)
+
+        max_mel_length = math.ceil(y.shape[-1] / self.hop_length)
+        spec = spec[..., :max_mel_length].transpose(1, 2)
+
+        if length is None:
+            return spec
+        else:
+            spec_len = torch.ceil(length / self.hop_length).clamp(max=spec.shape[1])
+            return spec, spec_len
+
+
+def length_to_mask(length, max_len=None, dtype=None, device=None):
+    """Creates a binary mask for each sequence.
+
+    Reference: https://discuss.pytorch.org/t/how-to-generate-variable-length-mask/23397/3
+
+    Arguments
+    ---------
+    length : torch.LongTensor
+        Containing the length of each sequence in the batch. Must be 1D.
+    max_len : int
+        Max length for the mask, also the size of the second dimension.
+    dtype : torch.dtype, default: None
+        The dtype of the generated mask.
+    device: torch.device, default: None
+        The device to put the mask variable.
+
+    Returns
+    -------
+    mask : tensor
+        The binary mask.
+
+    Example
+    -------
+    >>> length=torch.Tensor([1,2,3])
+    >>> mask=length_to_mask(length)
+    >>> mask
+    tensor([[1., 0., 0.],
+            [1., 1., 0.],
+            [1., 1., 1.]])
+    """
+    assert len(length.shape) == 1
+
+    if max_len is None:
+        max_len = length.max().long().item()  # using arange to generate mask
+    mask = torch.arange(max_len, device=length.device, dtype=length.dtype).expand(
+        len(length), max_len
+    ) < length.unsqueeze(1)
+
+    if dtype is None:
+        dtype = length.dtype
+
+    if device is None:
+        device = length.device
+
+    mask = torch.as_tensor(mask, dtype=dtype, device=device)
+    return mask
+
+
+def get_padding_elem(L_in: int, stride: int, kernel_size: int, dilation: int):
+    """This function computes the number of elements to add for zero-padding.
+
+    Arguments
+    ---------
+    L_in : int
+    stride: int
+    kernel_size : int
+    dilation : int
+    """
+    if stride > 1:
+        n_steps = math.ceil(((L_in - kernel_size * dilation) / stride) + 1)
+        L_out = stride * (n_steps - 1) + kernel_size * dilation
+        padding = [kernel_size // 2, kernel_size // 2]
+
+    else:
+        L_out = (L_in - dilation * (kernel_size - 1) - 1) // stride + 1
+
+        padding = [(L_in - L_out) // 2, (L_in - L_out) // 2]
+    return padding
+
+
+class Conv1d(nn.Module):
+    """This function implements 1d convolution.
+
+    Arguments
+    ---------
+    out_channels : int
+        It is the number of output channels.
+    kernel_size : int
+        Kernel size of the convolutional filters.
+    input_shape : tuple
+        The shape of the input. Alternatively use ``in_channels``.
+    in_channels : int
+        The number of input channels. Alternatively use ``input_shape``.
+    stride : int
+        Stride factor of the convolutional filters. When the stride factor > 1,
+        a decimation in time is performed.
+    dilation : int
+        Dilation factor of the convolutional filters.
+    padding : str
+        (same, valid, causal). If "valid", no padding is performed.
+        If "same" and stride is 1, output shape is the same as the input shape.
+        "causal" results in causal (dilated) convolutions.
+    padding_mode : str
+        This flag specifies the type of padding. See torch.nn documentation
+        for more information.
+    skip_transpose : bool
+        If False, uses batch x time x channel convention of speechbrain.
+        If True, uses batch x channel x time convention.
+
+    Example
+    -------
+    >>> inp_tensor = torch.rand([10, 40, 16])
+    >>> cnn_1d = Conv1d(
+    ...     input_shape=inp_tensor.shape, out_channels=8, kernel_size=5
+    ... )
+    >>> out_tensor = cnn_1d(inp_tensor)
+    >>> out_tensor.shape
+    torch.Size([10, 40, 8])
+    """
+
+    def __init__(
+        self,
+        out_channels,
+        kernel_size,
+        input_shape=None,
+        in_channels=None,
+        stride=1,
+        dilation=1,
+        padding="same",
+        groups=1,
+        bias=True,
+        padding_mode="reflect",
+        skip_transpose=True,
+    ):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.dilation = dilation
+        self.padding = padding
+        self.padding_mode = padding_mode
+        self.unsqueeze = False
+        self.skip_transpose = skip_transpose
+
+        if input_shape is None and in_channels is None:
+            raise ValueError("Must provide one of input_shape or in_channels")
+
+        if in_channels is None:
+            in_channels = self._check_input_shape(input_shape)
+
+        self.conv = nn.Conv1d(
+            in_channels,
+            out_channels,
+            self.kernel_size,
+            stride=self.stride,
+            dilation=self.dilation,
+            padding=0,
+            groups=groups,
+            bias=bias,
+        )
+
+    def forward(self, x):
+        """Returns the output of the convolution.
+
+        Arguments
+        ---------
+        x : torch.Tensor (batch, time, channel)
+            input to convolve. 2d or 4d tensors are expected.
+        """
+
+        if not self.skip_transpose:
+            x = x.transpose(1, -1)
+
+        if self.unsqueeze:
+            x = x.unsqueeze(1)
+
+        if self.padding == "same":
+            x = self._manage_padding(x, self.kernel_size, self.dilation, self.stride)
+
+        elif self.padding == "causal":
+            num_pad = (self.kernel_size - 1) * self.dilation
+            x = F.pad(x, (num_pad, 0))
+
+        elif self.padding == "valid":
+            pass
+
+        else:
+            raise ValueError(
+                "Padding must be 'same', 'valid' or 'causal'. Got " + self.padding
+            )
+
+        wx = self.conv(x)
+
+        if self.unsqueeze:
+            wx = wx.squeeze(1)
+
+        if not self.skip_transpose:
+            wx = wx.transpose(1, -1)
+
+        return wx
+
+    def _manage_padding(
+        self,
+        x,
+        kernel_size: int,
+        dilation: int,
+        stride: int,
+    ):
+        """This function performs zero-padding on the time axis
+        such that their lengths is unchanged after the convolution.
+
+        Arguments
+        ---------
+        x : torch.Tensor
+            Input tensor.
+        kernel_size : int
+            Size of kernel.
+        dilation : int
+            Dilation used.
+        stride : int
+            Stride.
+        """
+
+        # Detecting input shape
+        L_in = x.shape[-1]
+
+        # Time padding
+        padding = get_padding_elem(L_in, stride, kernel_size, dilation)
+
+        # Applying padding
+        x = F.pad(x, padding, mode=self.padding_mode)
+
+        return x
+
+    def _check_input_shape(self, shape):
+        """Checks the input shape and returns the number of input channels."""
+
+        if len(shape) == 2:
+            self.unsqueeze = True
+            in_channels = 1
+        elif self.skip_transpose:
+            in_channels = shape[1]
+        elif len(shape) == 3:
+            in_channels = shape[2]
+        else:
+            raise ValueError("conv1d expects 2d, 3d inputs. Got " + str(len(shape)))
+
+        # Kernel size must be odd
+        if self.kernel_size % 2 == 0:
+            raise ValueError(
+                "The field kernel size must be an odd number. Got %s."
+                % (self.kernel_size)
+            )
+        return in_channels
+
+
+class Fp32BatchNorm(nn.Module):
+    def __init__(self, sync=True, *args, **kwargs):
+        super().__init__()
+
+        if (
+            not torch.distributed.is_initialized()
+            or torch.distributed.get_world_size() == 1
+        ):
+            sync = False
+
+        if sync:
+            self.bn = nn.SyncBatchNorm(*args, **kwargs)
+        else:
+            self.bn = nn.BatchNorm1d(*args, **kwargs)
+
+        self.sync = sync
+
+    def forward(self, input):
+        if self.bn.running_mean.dtype != torch.float:
+            if self.sync:
+                self.bn.running_mean = self.bn.running_mean.float()
+                self.bn.running_var = self.bn.running_var.float()
+                if self.bn.affine:
+                    try:
+                        self.bn.weight = self.bn.weight.float()
+                        self.bn.bias = self.bn.bias.float()
+                    except:
+                        self.bn.float()
+            else:
+                self.bn.float()
+
+        output = self.bn(input.float())
+        return output.type_as(input)
+
+
+class BatchNorm1d(nn.Module):
+    """Applies 1d batch normalization to the input tensor.
+
+    Arguments
+    ---------
+    input_shape : tuple
+        The expected shape of the input. Alternatively, use ``input_size``.
+    input_size : int
+        The expected size of the input. Alternatively, use ``input_shape``.
+    eps : float
+        This value is added to std deviation estimation to improve the numerical
+        stability.
+    momentum : float
+        It is a value used for the running_mean and running_var computation.
+    affine : bool
+        When set to True, the affine parameters are learned.
+    track_running_stats : bool
+        When set to True, this module tracks the running mean and variance,
+        and when set to False, this module does not track such statistics.
+    combine_batch_time : bool
+        When true, it combines batch an time axis.
+
+
+    Example
+    -------
+    >>> input = torch.randn(100, 10)
+    >>> norm = BatchNorm1d(input_shape=input.shape)
+    >>> output = norm(input)
+    >>> output.shape
+    torch.Size([100, 10])
+    """
+
+    def __init__(
+        self,
+        input_shape=None,
+        input_size=None,
+        eps=1e-05,
+        momentum=0.1,
+        affine=True,
+        track_running_stats=True,
+        combine_batch_time=False,
+        skip_transpose=True,
+        enabled=True,
+    ):
+        super().__init__()
+        self.combine_batch_time = combine_batch_time
+        self.skip_transpose = skip_transpose
+
+        if input_size is None and skip_transpose:
+            input_size = input_shape[1]
+        elif input_size is None:
+            input_size = input_shape[-1]
+
+        if enabled:
+            self.norm = Fp32BatchNorm(
+                num_features=input_size,
+                eps=eps,
+                momentum=momentum,
+                affine=affine,
+                track_running_stats=track_running_stats,
+            )
+        else:
+            self.norm = nn.Identity()
+
+    def forward(self, x):
+        """Returns the normalized input tensor.
+
+        Arguments
+        ---------
+        x : torch.Tensor (batch, time, [channels])
+            input to normalize. 2d or 3d tensors are expected in input
+            4d tensors can be used when combine_dims=True.
+        """
+        shape_or = x.shape
+        if self.combine_batch_time:
+            if x.ndim == 3:
+                x = x.reshape(shape_or[0] * shape_or[1], shape_or[2])
+            else:
+                x = x.reshape(shape_or[0] * shape_or[1], shape_or[3], shape_or[2])
+
+        elif not self.skip_transpose:
+            x = x.transpose(-1, 1)
+
+        x_n = self.norm(x)
+
+        if self.combine_batch_time:
+            x_n = x_n.reshape(shape_or)
+        elif not self.skip_transpose:
+            x_n = x_n.transpose(1, -1)
+
+        return x_n
+
+
+class Linear(torch.nn.Module):
+    """Computes a linear transformation y = wx + b.
+
+    Arguments
+    ---------
+    n_neurons : int
+        It is the number of output neurons (i.e, the dimensionality of the
+        output).
+    bias : bool
+        If True, the additive bias b is adopted.
+    combine_dims : bool
+        If True and the input is 4D, combine 3rd and 4th dimensions of input.
+
+    Example
+    -------
+    >>> inputs = torch.rand(10, 50, 40)
+    >>> lin_t = Linear(input_shape=(10, 50, 40), n_neurons=100)
+    >>> output = lin_t(inputs)
+    >>> output.shape
+    torch.Size([10, 50, 100])
+    """
+
+    def __init__(
+        self,
+        n_neurons,
+        input_shape=None,
+        input_size=None,
+        bias=True,
+        combine_dims=False,
+    ):
+        super().__init__()
+        self.combine_dims = combine_dims
+
+        if input_shape is None and input_size is None:
+            raise ValueError("Expected one of input_shape or input_size")
+
+        if input_size is None:
+            input_size = input_shape[-1]
+            if len(input_shape) == 4 and self.combine_dims:
+                input_size = input_shape[2] * input_shape[3]
+
+        # Weights are initialized following pytorch approach
+        self.w = nn.Linear(input_size, n_neurons, bias=bias)
+
+    def forward(self, x):
+        """Returns the linear transformation of input tensor.
+
+        Arguments
+        ---------
+        x : torch.Tensor
+            Input to transform linearly.
+        """
+        if x.ndim == 4 and self.combine_dims:
+            x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3])
+
+        wx = self.w(x)
+
+        return wx
+
+
+class TDNNBlock(nn.Module):
+    """An implementation of TDNN.
+
+    Arguments
+    ----------
+    in_channels : int
+        Number of input channels.
+    out_channels : int
+        The number of output channels.
+    kernel_size : int
+        The kernel size of the TDNN blocks.
+    dilation : int
+        The dilation of the Res2Net block.
+    activation : torch class
+        A class for constructing the activation layers.
+
+    Example
+    -------
+    >>> inp_tensor = torch.rand([8, 120, 64]).transpose(1, 2)
+    >>> layer = TDNNBlock(64, 64, kernel_size=3, dilation=1)
+    >>> out_tensor = layer(inp_tensor).transpose(1, 2)
+    >>> out_tensor.shape
+    torch.Size([8, 120, 64])
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        dilation,
+        activation=nn.ReLU,
+        batch_norm=True,
+    ):
+        super(TDNNBlock, self).__init__()
+        self.conv = Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            dilation=dilation,
+        )
+        self.activation = activation()
+        self.norm = BatchNorm1d(input_size=out_channels, enabled=batch_norm)
+
+    def forward(self, x):
+        return self.norm(self.activation(self.conv(x)))
+
+
+class Res2NetBlock(torch.nn.Module):
+    """An implementation of Res2NetBlock w/ dilation.
+
+    Arguments
+    ---------
+    in_channels : int
+        The number of channels expected in the input.
+    out_channels : int
+        The number of output channels.
+    scale : int
+        The scale of the Res2Net block.
+    kernel_size: int
+        The kernel size of the Res2Net block.
+    dilation : int
+        The dilation of the Res2Net block.
+
+    Example
+    -------
+    >>> inp_tensor = torch.rand([8, 120, 64]).transpose(1, 2)
+    >>> layer = Res2NetBlock(64, 64, scale=4, dilation=3)
+    >>> out_tensor = layer(inp_tensor).transpose(1, 2)
+    >>> out_tensor.shape
+    torch.Size([8, 120, 64])
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        scale=8,
+        kernel_size=3,
+        dilation=1,
+        batch_norm=True,
+    ):
+        super(Res2NetBlock, self).__init__()
+        assert in_channels % scale == 0
+        assert out_channels % scale == 0
+
+        in_channel = in_channels // scale
+        hidden_channel = out_channels // scale
+
+        self.blocks = nn.ModuleList(
+            [
+                TDNNBlock(
+                    in_channel,
+                    hidden_channel,
+                    kernel_size=kernel_size,
+                    dilation=dilation,
+                    batch_norm=batch_norm,
+                )
+                for i in range(scale - 1)
+            ]
+        )
+        self.scale = scale
+
+    def forward(self, x):
+        y = []
+        for i, x_i in enumerate(torch.chunk(x, self.scale, dim=1)):
+            if i == 0:
+                y_i = x_i
+            elif i == 1:
+                y_i = self.blocks[i - 1](x_i)
+            else:
+                y_i = self.blocks[i - 1](x_i + y_i)
+            y.append(y_i)
+        y = torch.cat(y, dim=1)
+        return y
+
+
+class SEBlock(nn.Module):
+    """An implementation of squeeze-and-excitation block.
+
+    Arguments
+    ---------
+    in_channels : int
+        The number of input channels.
+    se_channels : int
+        The number of output channels after squeeze.
+    out_channels : int
+        The number of output channels.
+
+    Example
+    -------
+    >>> inp_tensor = torch.rand([8, 120, 64]).transpose(1, 2)
+    >>> se_layer = SEBlock(64, 16, 64)
+    >>> lengths = torch.rand((8,))
+    >>> out_tensor = se_layer(inp_tensor, lengths).transpose(1, 2)
+    >>> out_tensor.shape
+    torch.Size([8, 120, 64])
+    """
+
+    def __init__(self, in_channels, se_channels, out_channels):
+        super(SEBlock, self).__init__()
+
+        self.conv1 = Conv1d(
+            in_channels=in_channels, out_channels=se_channels, kernel_size=1
+        )
+        self.relu = torch.nn.ReLU(inplace=True)
+        self.conv2 = Conv1d(
+            in_channels=se_channels, out_channels=out_channels, kernel_size=1
+        )
+        self.sigmoid = torch.nn.Sigmoid()
+
+    def forward(self, x, lengths=None):
+        L = x.shape[-1]
+        if lengths is not None:
+            mask = length_to_mask(lengths * L, max_len=L, device=x.device)
+            mask = mask.unsqueeze(1)
+            total = mask.sum(dim=2, keepdim=True)
+            s = (x * mask).sum(dim=2, keepdim=True) / total
+        else:
+            s = x.mean(dim=2, keepdim=True)
+
+        s = self.relu(self.conv1(s))
+        s = self.sigmoid(self.conv2(s))
+
+        return s * x
+
+
+class AttentiveStatisticsPooling(nn.Module):
+    """This class implements an attentive statistic pooling layer for each channel.
+    It returns the concatenated mean and std of the input tensor.
+
+    Arguments
+    ---------
+    channels: int
+        The number of input channels.
+    attention_channels: int
+        The number of attention channels.
+
+    Example
+    -------
+    >>> inp_tensor = torch.rand([8, 120, 64]).transpose(1, 2)
+    >>> asp_layer = AttentiveStatisticsPooling(64)
+    >>> lengths = torch.rand((8,))
+    >>> out_tensor = asp_layer(inp_tensor, lengths).transpose(1, 2)
+    >>> out_tensor.shape
+    torch.Size([8, 1, 128])
+    """
+
+    def __init__(
+        self, channels, attention_channels=128, global_context=True, batch_norm=True
+    ):
+        super().__init__()
+
+        self.eps = 1e-12
+        self.global_context = global_context
+        if global_context:
+            self.tdnn = TDNNBlock(
+                channels * 3, attention_channels, 1, 1, batch_norm=batch_norm
+            )
+        else:
+            self.tdnn = TDNNBlock(
+                channels, attention_channels, 1, 1, batch_norm, batch_norm
+            )
+        self.tanh = nn.Tanh()
+        self.conv = Conv1d(
+            in_channels=attention_channels, out_channels=channels, kernel_size=1
+        )
+
+    def forward(self, x, lengths=None):
+        """Calculates mean and std for a batch (input tensor).
+
+        Arguments
+        ---------
+        x : torch.Tensor
+            Tensor of shape [N, C, L].
+        """
+        L = x.shape[-1]
+
+        def _compute_statistics(x, m, dim=2, eps=self.eps):
+            mean = (m * x).sum(dim)
+            std = torch.sqrt((m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(eps))
+            return mean, std
+
+        if lengths is None:
+            lengths = torch.ones(x.shape[0], device=x.device)
+
+        # Make binary mask of shape [N, 1, L]
+        mask = length_to_mask(lengths * L, max_len=L, device=x.device)
+        mask = mask.unsqueeze(1)
+
+        # Expand the temporal context of the pooling layer by allowing the
+        # self-attention to look at global properties of the utterance.
+        if self.global_context:
+            # torch.std is unstable for backward computation
+            # https://github.com/pytorch/pytorch/issues/4320
+            total = mask.sum(dim=2, keepdim=True).float()
+            mean, std = _compute_statistics(x, mask / total)
+            mean = mean.unsqueeze(2).repeat(1, 1, L)
+            std = std.unsqueeze(2).repeat(1, 1, L)
+            attn = torch.cat([x, mean, std], dim=1)
+        else:
+            attn = x
+
+        # Apply layers
+        attn = self.conv(self.tanh(self.tdnn(attn)))
+
+        # Filter out zero-paddings
+        attn = attn.masked_fill(mask == 0, float("-inf"))
+
+        attn = F.softmax(attn, dim=2)
+        mean, std = _compute_statistics(x, attn)
+        # Append mean and std of the batch
+        pooled_stats = torch.cat((mean, std), dim=1)
+        pooled_stats = pooled_stats.unsqueeze(2)
+
+        return pooled_stats
+
+
+class SERes2NetBlock(nn.Module):
+    """An implementation of building block in ECAPA-TDNN, i.e.,
+    TDNN-Res2Net-TDNN-SEBlock.
+
+    Arguments
+    ----------
+    out_channels: int
+        The number of output channels.
+    res2net_scale: int
+        The scale of the Res2Net block.
+    kernel_size: int
+        The kernel size of the TDNN blocks.
+    dilation: int
+        The dilation of the Res2Net block.
+    activation : torch class
+        A class for constructing the activation layers.
+
+    Example
+    -------
+    >>> x = torch.rand(8, 120, 64).transpose(1, 2)
+    >>> conv = SERes2NetBlock(64, 64, res2net_scale=4)
+    >>> out = conv(x).transpose(1, 2)
+    >>> out.shape
+    torch.Size([8, 120, 64])
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        res2net_scale=8,
+        se_channels=128,
+        kernel_size=1,
+        dilation=1,
+        activation=torch.nn.ReLU,
+        batch_norm=True,
+    ):
+        super().__init__()
+        self.out_channels = out_channels
+        self.tdnn1 = TDNNBlock(
+            in_channels,
+            out_channels,
+            kernel_size=1,
+            dilation=1,
+            activation=activation,
+            batch_norm=batch_norm,
+        )
+        self.res2net_block = Res2NetBlock(
+            out_channels,
+            out_channels,
+            res2net_scale,
+            kernel_size,
+            dilation,
+            batch_norm=batch_norm,
+        )
+        self.tdnn2 = TDNNBlock(
+            out_channels,
+            out_channels,
+            kernel_size=1,
+            dilation=1,
+            activation=activation,
+            batch_norm=batch_norm,
+        )
+        self.se_block = SEBlock(out_channels, se_channels, out_channels)
+
+        self.shortcut = None
+        if in_channels != out_channels:
+            self.shortcut = Conv1d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=1,
+            )
+
+    def forward(self, x, lengths=None):
+        residual = x
+        if self.shortcut:
+            residual = self.shortcut(x)
+
+        x = self.tdnn1(x)
+        x = self.res2net_block(x)
+        x = self.tdnn2(x)
+        x = self.se_block(x, lengths)
+
+        return x + residual
+
+
+class ECAPA_TDNN(torch.nn.Module):
+    """An implementation of the speaker embedding model in a paper.
+    "ECAPA-TDNN: Emphasized Channel Attention, Propagation and Aggregation in
+    TDNN Based Speaker Verification" (https://arxiv.org/abs/2005.07143).
+
+    Arguments
+    ---------
+    device : str
+        Device used, e.g., "cpu" or "cuda".
+    activation : torch class
+        A class for constructing the activation layers.
+    channels : list of ints
+        Output channels for TDNN/SERes2Net layer.
+    kernel_sizes : list of ints
+        List of kernel sizes for each layer.
+    dilations : list of ints
+        List of dilations for kernels in each layer.
+    lin_neurons : int
+        Number of neurons in linear layers.
+
+    Example
+    -------
+    >>> input_feats = torch.rand([5, 120, 80])
+    >>> compute_embedding = ECAPA_TDNN(80, lin_neurons=192)
+    >>> outputs = compute_embedding(input_feats)
+    >>> outputs.shape
+    torch.Size([5, 1, 192])
+    """
+
+    def __init__(
+        self,
+        input_size,
+        lin_neurons=192,
+        activation=torch.nn.ReLU,
+        channels=[512, 512, 512, 512, 1536],
+        kernel_sizes=[5, 3, 3, 3, 1],
+        dilations=[1, 2, 3, 4, 1],
+        attention_channels=128,
+        res2net_scale=8,
+        se_channels=128,
+        global_context=True,
+        batch_norm=True,
+    ):
+
+        super().__init__()
+        assert len(channels) == len(kernel_sizes)
+        assert len(channels) == len(dilations)
+        self.channels = channels
+        self.blocks = nn.ModuleList()
+
+        # The initial TDNN layer
+        self.blocks.append(
+            TDNNBlock(
+                input_size,
+                channels[0],
+                kernel_sizes[0],
+                dilations[0],
+                activation,
+                batch_norm=batch_norm,
+            )
+        )
+
+        # SE-Res2Net layers
+        for i in range(1, len(channels) - 1):
+            self.blocks.append(
+                SERes2NetBlock(
+                    channels[i - 1],
+                    channels[i],
+                    res2net_scale=res2net_scale,
+                    se_channels=se_channels,
+                    kernel_size=kernel_sizes[i],
+                    dilation=dilations[i],
+                    activation=activation,
+                    batch_norm=batch_norm,
+                )
+            )
+
+        # Multi-layer feature aggregation
+        self.mfa = TDNNBlock(
+            channels[-1],
+            channels[-1],
+            kernel_sizes[-1],
+            dilations[-1],
+            activation,
+            batch_norm=batch_norm,
+        )
+
+        # Attentive Statistical Pooling
+        self.asp = AttentiveStatisticsPooling(
+            channels[-1],
+            attention_channels=attention_channels,
+            global_context=global_context,
+            batch_norm=batch_norm,
+        )
+        self.asp_bn = BatchNorm1d(input_size=channels[-1] * 2, enabled=batch_norm)
+
+        # Final linear transformation
+        self.fc = Conv1d(
+            in_channels=channels[-1] * 2,
+            out_channels=lin_neurons,
+            kernel_size=1,
+        )
+
+    # @torch.cuda.amp.autocast(enabled=True, dtype=torch.float32)
+    def forward(self, x, lengths=None):
+        """Returns the embedding vector.
+
+        Arguments
+        ---------
+        x : torch.Tensor
+            Tensor of shape (batch, time, channel).
+        """
+        # Minimize transpose for efficiency
+        x = x.transpose(1, 2)
+
+        xl = []
+        for layer in self.blocks:
+            try:
+                x = layer(x, lengths=lengths)
+            except TypeError:
+                x = layer(x)
+            xl.append(x)
+
+        # Multi-layer feature aggregation
+        x = torch.cat(xl[1:], dim=1)
+        x = self.mfa(x)
+
+        # Attentive Statistical Pooling
+        x = self.asp(x, lengths=lengths)
+        x = self.asp_bn(x)
+
+        # Final linear transformation
+        x = self.fc(x)
+
+        x = x.squeeze(-1)
+        return x
+
+
+class SpeakerEmbedddingExtractor(object):
+
+    def __init__(self, ckpt_path, device="cuda"):
+        # NOTE: The sampling rate is 16000
+        self.mel_extractor = TorchMelSpectrogram()
+        self.mel_extractor.to(device)
+        model = ECAPA_TDNN(
+            80,
+            512,
+            channels=[512, 512, 512, 512, 1536],
+            kernel_sizes=[5, 3, 3, 3, 1],
+            dilations=[1, 2, 3, 4, 1],
+            attention_channels=128,
+            res2net_scale=4,
+            se_channels=128,
+            global_context=True,
+            batch_norm=True,
+        )
+        model.load_state_dict(torch.load(ckpt_path), strict=True)
+        model.eval()
+        self.model = model
+        self.model.to(device)
+
+    def __call__(self, wav):
+        # wav, sr = torchaudio.load(audio_path)
+        # assert sr == 16000, f"The sampling rate is not 16000"
+        # print(wav.shape)
+        mel = self.mel_extractor(wav.unsqueeze(0))
+        spk = self.model(mel)
+        spk = spk[0]
+        return spk

fireredtts/modules/flow/__init__.py

@@ -0,0 +1,24 @@
+from fireredtts.modules.flow.codec_embedding import HHGCodecEmbedding
+from fireredtts.modules.flow.conformer import ConformerDecoderV2
+from fireredtts.modules.flow.mel_encoder import MelReduceEncoder
+from fireredtts.modules.flow.decoder import ConditionalCFM, ConditionalDecoder
+from fireredtts.modules.flow.flow_model import InterpolateRegulator, CrossAttnFlowMatching
+from fireredtts.modules.flow.mel_spectrogram import MelSpectrogramExtractor
+
+
+def get_flow_frontend(flow_config):
+    flow = CrossAttnFlowMatching(
+        output_size=flow_config["output_size"],
+        input_embedding=HHGCodecEmbedding(**flow_config["input_embedding"]),
+        encoder=ConformerDecoderV2(**flow_config["encoder"]),
+        length_regulator=InterpolateRegulator(**flow_config["length_regulator"]),
+        mel_encoder=MelReduceEncoder(**flow_config["mel_encoder"]),
+        decoder=ConditionalCFM(
+            estimator=ConditionalDecoder(**flow_config["decoder"]["estimator"]),
+            t_scheduler=flow_config["decoder"]["t_scheduler"],
+            inference_cfg_rate=flow_config["decoder"]["inference_cfg_rate"]
+        )
+    )
+    return flow
+
+

fireredtts/modules/flow/codebook.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f32d6280c561e4d0bf741e3ea07d651aee71da84c17fc55b686cb825bd677656
+size 131200

fireredtts/modules/flow/codec_embedding.py

@@ -0,0 +1,30 @@
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+class HHGCodecEmbedding(nn.Module):
+    def __init__(self, out_channels, codebook_path:str, freeze=True):
+        super().__init__()
+        # (2, 128, 128)
+        codebook = torch.from_numpy(np.load(codebook_path).copy())
+        assert codebook.shape[0] == 2 and codebook.shape[1] == 128
+        self.codebook_dim = codebook.shape[2]
+
+        self.codebook = torch.nn.ModuleList([
+            torch.nn.Embedding.from_pretrained(codebook[i], freeze=freeze) 
+            for i in range(codebook.shape[0])]
+        )
+        if self.codebook_dim * 2 != out_channels:
+            self.proj = nn.Linear(self.codebook_dim * 2, out_channels)
+        else:
+            self.proj = nn.Identity()
+
+    def forward(self, tokens):
+        token_embs = torch.cat([
+            self.codebook[0](tokens % 128),
+            self.codebook[1](tokens // 128)
+        ], dim=-1)
+        token_embs = self.proj(token_embs)
+        return token_embs
+

fireredtts/modules/flow/conformer.py

@@ -0,0 +1,730 @@
+import typing as tp
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fireredtts.modules.flow.utils import make_pad_mask
+
+
+class MultiHeadedAttention(nn.Module):
+    """Multi-Head Attention layer.
+
+    Args:
+        n_head (int): The number of heads.
+        n_feat (int): The number of features.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(self,
+                 n_head: int,
+                 n_feat: int,
+                 dropout_rate: float,
+                 key_bias: bool = True):
+        """Construct an MultiHeadedAttention object."""
+        super().__init__()
+        assert n_feat % n_head == 0
+        # We assume d_v always equals d_k
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        self.linear_q = nn.Linear(n_feat, n_feat)
+        self.linear_k = nn.Linear(n_feat, n_feat, bias=key_bias)
+        self.linear_v = nn.Linear(n_feat, n_feat)
+        self.linear_out = nn.Linear(n_feat, n_feat)
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+    def forward_qkv(
+        self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
+    ) -> tp.Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Transform query, key and value.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+
+        Returns:
+            torch.Tensor: Transformed query tensor, size
+                (#batch, n_head, time1, d_k).
+            torch.Tensor: Transformed key tensor, size
+                (#batch, n_head, time2, d_k).
+            torch.Tensor: Transformed value tensor, size
+                (#batch, n_head, time2, d_k).
+
+        """
+        n_batch = query.size(0)
+        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
+        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
+        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
+        q = q.transpose(1, 2)  # (batch, head, time1, d_k)
+        k = k.transpose(1, 2)  # (batch, head, time2, d_k)
+        v = v.transpose(1, 2)  # (batch, head, time2, d_k)
+
+        return q, k, v
+
+    def forward_attention(
+        self,
+        value: torch.Tensor,
+        scores: torch.Tensor,
+        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool)
+    ) -> torch.Tensor:
+        """Compute attention context vector.
+
+        Args:
+            value (torch.Tensor): Transformed value, size
+                (#batch, n_head, time2, d_k).
+            scores (torch.Tensor): Attention score, size
+                (#batch, n_head, time1, time2).
+            mask (torch.Tensor): Mask, size (#batch, 1, time2) or
+                (#batch, time1, time2), (0, 0, 0) means fake mask.
+
+        Returns:
+            torch.Tensor: Transformed value (#batch, time1, d_model)
+                weighted by the attention score (#batch, time1, time2).
+
+        """
+        n_batch = value.size(0)
+        # NOTE(xcsong): When will `if mask.size(2) > 0` be True?
+        #   1. onnx(16/4) [WHY? Because we feed real cache & real mask for the
+        #           1st chunk to ease the onnx export.]
+        #   2. pytorch training
+        if mask.size(2) > 0:  # time2 > 0
+            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
+            # For last chunk, time2 might be larger than scores.size(-1)
+            mask = mask[:, :, :, :scores.size(-1)]  # (batch, 1, *, time2)
+            scores = scores.masked_fill(mask, -float('inf'))
+            attn = torch.softmax(scores, dim=-1).masked_fill(
+                mask, 0.0)  # (batch, head, time1, time2)
+        # NOTE(xcsong): When will `if mask.size(2) > 0` be False?
+        #   1. onnx(16/-1, -1/-1, 16/0)
+        #   2. jit (16/-1, -1/-1, 16/0, 16/4)
+        else:
+            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
+
+        p_attn = self.dropout(attn)
+        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
+        x = (x.transpose(1, 2).contiguous().view(n_batch, -1,
+                                                 self.h * self.d_k)
+             )  # (batch, time1, d_model)
+
+        return self.linear_out(x)  # (batch, time1, d_model)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+        pos_emb: torch.Tensor = torch.empty(0),
+        cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
+    ) -> tp.Tuple[torch.Tensor, torch.Tensor]:
+        """Compute scaled dot product attention.
+
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
+                (#batch, time1, time2).
+                1.When applying cross attention between decoder and encoder,
+                the batch padding mask for input is in (#batch, 1, T) shape.
+                2.When applying self attention of encoder,
+                the mask is in (#batch, T, T)  shape.
+            cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2),
+                where `cache_t == chunk_size * num_decoding_left_chunks`
+                and `head * d_k == size`
+
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time1, d_model).
+            torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2)
+                where `cache_t == chunk_size * num_decoding_left_chunks`
+                and `head * d_k == size`
+
+        """
+        q, k, v = self.forward_qkv(query, key, value)
+
+        # NOTE(xcsong):
+        #   when export onnx model, for 1st chunk, we feed
+        #       cache(1, head, 0, d_k * 2) (16/-1, -1/-1, 16/0 mode)
+        #       or cache(1, head, real_cache_t, d_k * 2) (16/4 mode).
+        #       In all modes, `if cache.size(0) > 0` will alwayse be `True`
+        #       and we will always do splitting and
+        #       concatnation(this will simplify onnx export). Note that
+        #       it's OK to concat & split zero-shaped tensors(see code below).
+        #   when export jit  model, for 1st chunk, we always feed
+        #       cache(0, 0, 0, 0) since jit supports dynamic if-branch.
+        # >>> a = torch.ones((1, 2, 0, 4))
+        # >>> b = torch.ones((1, 2, 3, 4))
+        # >>> c = torch.cat((a, b), dim=2)
+        # >>> torch.equal(b, c)        # True
+        # >>> d = torch.split(a, 2, dim=-1)
+        # >>> torch.equal(d[0], d[1])  # True
+        if cache.size(0) > 0:
+            key_cache, value_cache = torch.split(cache,
+                                                 cache.size(-1) // 2,
+                                                 dim=-1)
+            k = torch.cat([key_cache, k], dim=2)
+            v = torch.cat([value_cache, v], dim=2)
+        # NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's
+        #   non-trivial to calculate `next_cache_start` here.
+        new_cache = torch.cat((k, v), dim=-1)
+
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+        return self.forward_attention(v, scores, mask), new_cache
+
+
+class RelPositionMultiHeadedAttention(MultiHeadedAttention):
+    """Multi-Head Attention layer with relative position encoding.
+    Paper: https://arxiv.org/abs/1901.02860
+    Args:
+        n_head (int): The number of heads.
+        n_feat (int): The number of features.
+        dropout_rate (float): Dropout rate.
+    """
+
+    def __init__(self,
+                 n_head: int,
+                 n_feat: int,
+                 dropout_rate: float,
+                 key_bias: bool = True):
+        """Construct an RelPositionMultiHeadedAttention object."""
+        super().__init__(n_head, n_feat, dropout_rate, key_bias)
+        # linear transformation for positional encoding
+        self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)
+        # these two learnable bias are used in matrix c and matrix d
+        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+        self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k))
+        self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k))
+        torch.nn.init.xavier_uniform_(self.pos_bias_u)
+        torch.nn.init.xavier_uniform_(self.pos_bias_v)
+
+    def rel_shift(self, x):
+        """Compute relative positional encoding.
+
+        Args:
+            x (torch.Tensor): Input tensor (batch, head, time1, 2*time1-1).
+            time1 means the length of query vector.
+
+        Returns:
+            torch.Tensor: Output tensor.
+
+        """
+        zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype)
+        x_padded = torch.cat([zero_pad, x], dim=-1)
+
+        x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2))
+        x = x_padded[:, :, 1:].view_as(x)[
+            :, :, :, : x.size(-1) // 2 + 1
+        ]  # only keep the positions from 0 to time2
+        return x
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+        pos_emb: torch.Tensor = torch.empty(0),
+        cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
+    ) -> tp.Tuple[torch.Tensor, torch.Tensor]:
+        """Compute 'Scaled Dot Product Attention' with rel. positional encoding.
+        Args:
+            query (torch.Tensor): Query tensor (#batch, time1, size).
+            key (torch.Tensor): Key tensor (#batch, time2, size).
+            value (torch.Tensor): Value tensor (#batch, time2, size).
+            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
+                (#batch, time1, time2), (0, 0, 0) means fake mask.
+            pos_emb (torch.Tensor): Positional embedding tensor
+                (#batch, time2, size).
+            cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2),
+                where `cache_t == chunk_size * num_decoding_left_chunks`
+                and `head * d_k == size`
+        Returns:
+            torch.Tensor: Output tensor (#batch, time1, d_model).
+            torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2)
+                where `cache_t == chunk_size * num_decoding_left_chunks`
+                and `head * d_k == size`
+        """
+        q, k, v = self.forward_qkv(query, key, value)
+        q = q.transpose(1, 2)  # (batch, time1, head, d_k)
+
+        # NOTE(xcsong):
+        #   when export onnx model, for 1st chunk, we feed
+        #       cache(1, head, 0, d_k * 2) (16/-1, -1/-1, 16/0 mode)
+        #       or cache(1, head, real_cache_t, d_k * 2) (16/4 mode).
+        #       In all modes, `if cache.size(0) > 0` will alwayse be `True`
+        #       and we will always do splitting and
+        #       concatnation(this will simplify onnx export). Note that
+        #       it's OK to concat & split zero-shaped tensors(see code below).
+        #   when export jit  model, for 1st chunk, we always feed
+        #       cache(0, 0, 0, 0) since jit supports dynamic if-branch.
+        # >>> a = torch.ones((1, 2, 0, 4))
+        # >>> b = torch.ones((1, 2, 3, 4))
+        # >>> c = torch.cat((a, b), dim=2)
+        # >>> torch.equal(b, c)        # True
+        # >>> d = torch.split(a, 2, dim=-1)
+        # >>> torch.equal(d[0], d[1])  # True
+        if cache.size(0) > 0:
+            key_cache, value_cache = torch.split(cache,
+                                                 cache.size(-1) // 2,
+                                                 dim=-1)
+            k = torch.cat([key_cache, k], dim=2)
+            v = torch.cat([value_cache, v], dim=2)
+        # NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's
+        #   non-trivial to calculate `next_cache_start` here.
+        new_cache = torch.cat((k, v), dim=-1)
+
+        n_batch_pos = pos_emb.size(0)
+        p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)
+        p = p.transpose(1, 2)  # (batch, head, time1, d_k)
+
+        # (batch, head, time1, d_k)
+        q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
+        # (batch, head, time1, d_k)
+        q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
+
+        # compute attention score
+        # first compute matrix a and matrix c
+        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
+        # (batch, head, time1, time2)
+        matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
+
+        # compute matrix b and matrix d
+        # (batch, head, time1, time2)
+        matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
+        # NOTE(Xiang Lyu): Keep rel_shift since espnet rel_pos_emb is used
+        if matrix_ac.shape != matrix_bd.shape:
+            matrix_bd = self.rel_shift(matrix_bd)
+
+        scores = (matrix_ac + matrix_bd) / math.sqrt(
+            self.d_k)  # (batch, head, time1, time2)
+
+        return self.forward_attention(v, scores, mask), new_cache
+
+
+class PositionwiseFeedForward(torch.nn.Module):
+    """Positionwise feed forward layer.
+
+    FeedForward are appied on each position of the sequence.
+    The output dim is same with the input dim.
+
+    Args:
+        idim (int): Input dimenstion.
+        hidden_units (int): The number of hidden units.
+        dropout_rate (float): Dropout rate.
+        activation (torch.nn.Module): Activation function
+    """
+
+    def __init__(
+            self,
+            idim: int,
+            hidden_units: int,
+            dropout_rate: float,
+            activation: torch.nn.Module = torch.nn.ReLU(),
+    ):
+        """Construct a PositionwiseFeedForward object."""
+        super(PositionwiseFeedForward, self).__init__()
+        self.w_1 = torch.nn.Linear(idim, hidden_units)
+        self.activation = activation
+        self.dropout = torch.nn.Dropout(dropout_rate)
+        self.w_2 = torch.nn.Linear(hidden_units, idim)
+
+    def forward(self, xs: torch.Tensor) -> torch.Tensor:
+        """Forward function.
+
+        Args:
+            xs: input tensor (B, L, D)
+        Returns:
+            output tensor, (B, L, D)
+        """
+        return self.w_2(self.dropout(self.activation(self.w_1(xs))))
+
+
+class ConformerDecoderLayer(nn.Module):
+    """Encoder layer module.
+    Args:
+        size (int): Input dimension.
+        self_attn (torch.nn.Module): Self-attention module instance.
+            `MultiHeadedAttention` or `RelPositionMultiHeadedAttention`
+            instance can be used as the argument.
+        src_attn (torch.nn.Module): Cross-attention module instance.
+            `MultiHeadedAttention` or `RelPositionMultiHeadedAttention`
+            instance can be used as the argument.
+        feed_forward (torch.nn.Module): Feed-forward module instance.
+            `PositionwiseFeedForward` instance can be used as the argument.
+        feed_forward_macaron (torch.nn.Module): Additional feed-forward module
+             instance.
+            `PositionwiseFeedForward` instance can be used as the argument.
+        conv_module (torch.nn.Module): Convolution module instance.
+            `ConvlutionModule` instance can be used as the argument.
+        dropout_rate (float): Dropout rate.
+        normalize_before (bool):
+            True: use layer_norm before each sub-block.
+            False: use layer_norm after each sub-block.
+    """
+
+    def __init__(
+        self,
+        size: int,
+        self_attn: torch.nn.Module,
+        src_attn: tp.Optional[torch.nn.Module] = None,
+        feed_forward: tp.Optional[nn.Module] = None,
+        feed_forward_macaron: tp.Optional[nn.Module] = None,
+        conv_module: tp.Optional[nn.Module] = None,
+        dropout_rate: float = 0.1,
+        normalize_before: bool = True,
+    ):
+        """Construct an EncoderLayer object."""
+        super().__init__()
+        self.self_attn = self_attn
+        self.src_attn = src_attn
+        self.feed_forward = feed_forward
+        self.feed_forward_macaron = feed_forward_macaron
+        self.conv_module = conv_module
+        self.norm_ff = nn.LayerNorm(size, eps=1e-5)  # for the FNN module
+        self.norm_mha = nn.LayerNorm(size, eps=1e-5)  # for the MHA module
+        if src_attn is not None:
+            self.norm_mha2 = nn.LayerNorm(size, eps=1e-5)  # for the MHA module(src_attn)
+        if feed_forward_macaron is not None:
+            self.norm_ff_macaron = nn.LayerNorm(size, eps=1e-5)
+            self.ff_scale = 0.5
+        else:
+            self.ff_scale = 1.0
+        if self.conv_module is not None:
+            self.norm_conv = nn.LayerNorm(size, eps=1e-5)  # for the CNN module
+            self.norm_final = nn.LayerNorm(
+                size, eps=1e-5)  # for the final output of the block
+        self.dropout = nn.Dropout(dropout_rate)
+        self.size = size
+        self.normalize_before = normalize_before
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor,
+        # src-attention
+        memory: torch.Tensor,
+        memory_mask: torch.Tensor,
+        pos_emb: torch.Tensor,
+        mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+        att_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
+        cnn_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
+    ) -> tp.Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Compute encoded features.
+
+        Args:
+            x (torch.Tensor): (#batch, time, size)
+            mask (torch.Tensor): Mask tensor for the input (#batch, time，time),
+                (0, 0, 0) means fake mask.
+            pos_emb (torch.Tensor): positional encoding, must not be None
+                for ConformerEncoderLayer.
+            mask_pad (torch.Tensor): batch padding mask used for conv module.
+                (#batch, 1, time), (0, 0, 0) means fake mask.
+            att_cache (torch.Tensor): Cache tensor of the KEY & VALUE
+                (#batch=1, head, cache_t1, d_k * 2), head * d_k == size.
+            cnn_cache (torch.Tensor): Convolution cache in conformer layer
+                (#batch=1, size, cache_t2)
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, size).
+            torch.Tensor: Mask tensor (#batch, time, time).
+            torch.Tensor: att_cache tensor,
+                (#batch=1, head, cache_t1 + time, d_k * 2).
+            torch.Tensor: cnn_cahce tensor (#batch, size, cache_t2).
+        """
+
+        # whether to use macaron style
+        if self.feed_forward_macaron is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm_ff_macaron(x)
+            x = residual + self.ff_scale * self.dropout(
+                self.feed_forward_macaron(x))
+            if not self.normalize_before:
+                x = self.norm_ff_macaron(x)
+
+        # multi-headed self-attention module
+        residual = x
+        if self.normalize_before:
+            x = self.norm_mha(x)
+        x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb,
+                                              att_cache)
+        x = residual + self.dropout(x_att)
+        if not self.normalize_before:
+            x = self.norm_mha(x)
+        
+        # multi-headed cross-attention module
+        if self.src_attn is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm_mha2(x)
+            x_att, _ = self.src_attn(x, memory, memory, memory_mask)
+            x = residual + self.dropout(x_att)
+            if not self.normalize_before:
+                x = self.norm_mha2(x)
+
+        # convolution module
+        # Fake new cnn cache here, and then change it in conv_module
+        new_cnn_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device)
+        if self.conv_module is not None:
+            residual = x
+            if self.normalize_before:
+                x = self.norm_conv(x)
+            x, new_cnn_cache = self.conv_module(x, mask_pad, cnn_cache)
+            x = residual + self.dropout(x)
+
+            if not self.normalize_before:
+                x = self.norm_conv(x)
+
+        # feed forward module
+        residual = x
+        if self.normalize_before:
+            x = self.norm_ff(x)
+
+        x = residual + self.ff_scale * self.dropout(self.feed_forward(x))
+        if not self.normalize_before:
+            x = self.norm_ff(x)
+
+        if self.conv_module is not None:
+            x = self.norm_final(x)
+
+        return x, mask, new_att_cache, new_cnn_cache
+
+
+class EspnetRelPositionalEncoding(torch.nn.Module):
+    """Relative positional encoding module (new implementation).
+
+    Details can be found in https://github.com/espnet/espnet/pull/2816.
+
+    See : Appendix B in https://arxiv.org/abs/1901.02860
+
+    Args:
+        d_model (int): Embedding dimension.
+        dropout_rate (float): Dropout rate.
+        max_len (int): Maximum input length.
+
+    """
+
+    def __init__(self, d_model, dropout_rate, max_len=5000):
+        """Construct an PositionalEncoding object."""
+        super(EspnetRelPositionalEncoding, self).__init__()
+        self.d_model = d_model
+        self.xscale = math.sqrt(self.d_model)
+        self.dropout = torch.nn.Dropout(p=dropout_rate)
+        self.pe = None
+        self.extend_pe(torch.tensor(0.0).expand(1, max_len))
+
+    def extend_pe(self, x):
+        """Reset the positional encodings."""
+        if self.pe is not None:
+            # self.pe contains both positive and negative parts
+            # the length of self.pe is 2 * input_len - 1
+            if self.pe.size(1) >= x.size(1) * 2 - 1:
+                if self.pe.dtype != x.dtype or self.pe.device != x.device:
+                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)
+                return
+        # Suppose `i` means to the position of query vecotr and `j` means the
+        # position of key vector. We use position relative positions when keys
+        # are to the left (i>j) and negative relative positions otherwise (i<j).
+        pe_positive = torch.zeros(x.size(1), self.d_model)
+        pe_negative = torch.zeros(x.size(1), self.d_model)
+        position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, self.d_model, 2, dtype=torch.float32)
+            * -(math.log(10000.0) / self.d_model)
+        )
+        pe_positive[:, 0::2] = torch.sin(position * div_term)
+        pe_positive[:, 1::2] = torch.cos(position * div_term)
+        pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)
+        pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)
+
+        # Reserve the order of positive indices and concat both positive and
+        # negative indices. This is used to support the shifting trick
+        # as in https://arxiv.org/abs/1901.02860
+        pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)
+        pe_negative = pe_negative[1:].unsqueeze(0)
+        pe = torch.cat([pe_positive, pe_negative], dim=1)
+        self.pe = pe.to(device=x.device, dtype=x.dtype)
+
+    def forward(self, x: torch.Tensor, offset: tp.Union[int, torch.Tensor] = 0):
+        """Add positional encoding.
+
+        Args:
+            x (torch.Tensor): Input tensor (batch, time, `*`).
+
+        Returns:
+            torch.Tensor: Encoded tensor (batch, time, `*`).
+
+        """
+        self.extend_pe(x)
+        x = x * self.xscale
+        pos_emb = self.position_encoding(size=x.size(1), offset=offset)
+        return self.dropout(x), self.dropout(pos_emb)
+
+    def position_encoding(self,
+                          offset: tp.Union[int, torch.Tensor],
+                          size: int) -> torch.Tensor:
+        """ For getting encoding in a streaming fashion
+
+        Attention!!!!!
+        we apply dropout only once at the whole utterance level in a none
+        streaming way, but will call this function several times with
+        increasing input size in a streaming scenario, so the dropout will
+        be applied several times.
+
+        Args:
+            offset (int or torch.tensor): start offset
+            size (int): required size of position encoding
+
+        Returns:
+            torch.Tensor: Corresponding encoding
+        """
+        pos_emb = self.pe[
+            :,
+            self.pe.size(1) // 2 - size + 1 : self.pe.size(1) // 2 + size,
+        ]
+        return pos_emb
+
+
+class LinearNoSubsampling(torch.nn.Module):
+    """Linear transform the input without subsampling
+
+    Args:
+        idim (int): Input dimension.
+        odim (int): Output dimension.
+        dropout_rate (float): Dropout rate.
+
+    """
+
+    def __init__(self, idim: int, odim: int, dropout_rate: float,
+                 pos_enc_class: torch.nn.Module):
+        """Construct an linear object."""
+        super().__init__()
+        self.out = torch.nn.Sequential(
+            torch.nn.Linear(idim, odim),
+            torch.nn.LayerNorm(odim, eps=1e-5),
+            torch.nn.Dropout(dropout_rate),
+        )
+        self.pos_enc = pos_enc_class
+        self.right_context = 0
+        self.subsampling_rate = 1
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor,
+        offset: tp.Union[int, torch.Tensor] = 0
+    ) -> tp.Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Input x.
+
+        Args:
+            x (torch.Tensor): Input tensor (#batch, time, idim).
+            x_mask (torch.Tensor): Input mask (#batch, 1, time).
+
+        Returns:
+            torch.Tensor: linear input tensor (#batch, time', odim),
+                where time' = time .
+            torch.Tensor: linear input mask (#batch, 1, time'),
+                where time' = time .
+
+        """
+        x = self.out(x)
+        x, pos_emb = self.pos_enc(x, offset)
+        return x, pos_emb, x_mask
+
+
+class ConformerDecoderV2(nn.Module):
+    def __init__(self,
+                 input_size: int = 512,
+                 output_size: int = 512,
+                 attention_heads: int = 8,
+                 linear_units: int = 2048,
+                 num_blocks: int = 6,
+                 dropout_rate: float = 0.01,
+                 srcattention_start_index: int = 0,
+                 srcattention_end_index: int = 2,
+                 attention_dropout_rate: float = 0.01,
+                 positional_dropout_rate: float = 0.01,
+                 key_bias: bool = True,
+                 normalize_before: bool = True,
+                 ):
+        super().__init__()
+        self.num_blocks = num_blocks
+        self.normalize_before = normalize_before
+        self.output_size = output_size
+
+        self.embed = LinearNoSubsampling(
+            input_size, 
+            output_size, 
+            dropout_rate, 
+            EspnetRelPositionalEncoding(output_size, positional_dropout_rate),
+        )
+
+        self.encoders = torch.nn.ModuleList()
+        for i in range(self.num_blocks):
+            # construct src attention
+            if srcattention_start_index <= i <= srcattention_end_index:
+                srcattention_layer = MultiHeadedAttention(
+                    attention_heads, 
+                    output_size, 
+                    attention_dropout_rate, 
+                    key_bias
+                )
+            else:
+                srcattention_layer = None
+            # construct self attention
+            selfattention_layer = RelPositionMultiHeadedAttention(
+                attention_heads, 
+                output_size, 
+                attention_dropout_rate, 
+                key_bias
+            )
+            # construct ffn
+            ffn_layer = PositionwiseFeedForward(
+                output_size, 
+                linear_units, 
+                dropout_rate, 
+                torch.nn.SiLU()
+            )
+            self.encoders.append(
+                ConformerDecoderLayer(
+                    output_size, 
+                    selfattention_layer, 
+                    srcattention_layer, 
+                    ffn_layer, 
+                    None, 
+                    None, 
+                    dropout_rate, 
+                    normalize_before=normalize_before
+                )
+            )
+        self.after_norm = torch.nn.LayerNorm(output_size, eps=1e-5)
+
+    def forward_layers(self, xs: torch.Tensor, chunk_masks: torch.Tensor,
+                       memory: torch.Tensor, memory_masks: torch.Tensor,
+                       pos_emb: torch.Tensor, mask_pad: torch.Tensor) -> torch.Tensor:
+        for layer in self.encoders:
+            xs, chunk_masks, _, _ = layer(xs, chunk_masks, memory, memory_masks, pos_emb, mask_pad)
+        return xs
+
+    def forward(self, 
+                xs:torch.Tensor, 
+                xs_lens:torch.Tensor, 
+                memory:torch.Tensor, 
+                memory_lens: torch.Tensor,
+                ):
+        T = xs.size(1)
+        masks = ~make_pad_mask(xs_lens, T).unsqueeze(1)  # (B, 1, T)
+        T2 = memory.size(1)
+        memory_masks = ~make_pad_mask(memory_lens, T2).unsqueeze(1) # (B, 1, T2)
+
+        xs, pos_emb, masks = self.embed(xs, masks)
+
+        xs = self.forward_layers(xs, masks, memory, memory_masks, pos_emb, masks)
+
+        if self.normalize_before:
+            xs = self.after_norm(xs)
+        
+        return xs, masks
+

fireredtts/modules/flow/decoder.py

@@ -0,0 +1,396 @@
+import math
+from typing import Optional
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import pack, rearrange, repeat
+from diffusers.models.activations import get_activation
+
+from fireredtts.modules.flow.transformer import BasicTransformerBlock
+
+
+class SinusoidalPosEmb(torch.nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+        assert self.dim % 2 == 0, "SinusoidalPosEmb requires dim to be even"
+
+    def forward(self, x, scale=1000):
+        if x.ndim < 1:
+            x = x.unsqueeze(0)
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
+        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+class Block1D(torch.nn.Module):
+    def __init__(self, dim, dim_out, groups=8):
+        super().__init__()
+        self.block = torch.nn.Sequential(
+            torch.nn.Conv1d(dim, dim_out, 3, padding=1),
+            torch.nn.GroupNorm(groups, dim_out),
+            nn.Mish(),
+        )
+
+    def forward(self, x, mask):
+        output = self.block(x * mask)
+        return output * mask
+
+
+class ResnetBlock1D(torch.nn.Module):
+    def __init__(self, dim, dim_out, time_emb_dim, groups=8):
+        super().__init__()
+        self.mlp = torch.nn.Sequential(nn.Mish(), torch.nn.Linear(time_emb_dim, dim_out))
+
+        self.block1 = Block1D(dim, dim_out, groups=groups)
+        self.block2 = Block1D(dim_out, dim_out, groups=groups)
+
+        self.res_conv = torch.nn.Conv1d(dim, dim_out, 1)
+
+    def forward(self, x, mask, time_emb):
+        h = self.block1(x, mask)
+        h += self.mlp(time_emb).unsqueeze(-1)
+        h = self.block2(h, mask)
+        output = h + self.res_conv(x * mask)
+        return output
+
+
+class Downsample1D(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.conv = torch.nn.Conv1d(dim, dim, 3, 2, 1)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class TimestepEmbedding(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        time_embed_dim: int,
+        act_fn: str = "silu",
+        out_dim: int = None,
+        post_act_fn: Optional[str] = None,
+        cond_proj_dim=None,
+    ):
+        super().__init__()
+
+        self.linear_1 = nn.Linear(in_channels, time_embed_dim)
+
+        if cond_proj_dim is not None:
+            self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False)
+        else:
+            self.cond_proj = None
+
+        self.act = get_activation(act_fn)
+
+        if out_dim is not None:
+            time_embed_dim_out = out_dim
+        else:
+            time_embed_dim_out = time_embed_dim
+        self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out)
+
+        if post_act_fn is None:
+            self.post_act = None
+        else:
+            self.post_act = get_activation(post_act_fn)
+
+    def forward(self, sample, condition=None):
+        if condition is not None:
+            sample = sample + self.cond_proj(condition)
+        sample = self.linear_1(sample)
+
+        if self.act is not None:
+            sample = self.act(sample)
+
+        sample = self.linear_2(sample)
+
+        if self.post_act is not None:
+            sample = self.post_act(sample)
+        return sample
+
+
+class Upsample1D(nn.Module):
+    """A 1D upsampling layer with an optional convolution.
+
+    Parameters:
+        channels (`int`):
+            number of channels in the inputs and outputs.
+        use_conv (`bool`, default `False`):
+            option to use a convolution.
+        use_conv_transpose (`bool`, default `False`):
+            option to use a convolution transpose.
+        out_channels (`int`, optional):
+            number of output channels. Defaults to `channels`.
+    """
+
+    def __init__(self, channels, use_conv=False, use_conv_transpose=True, out_channels=None, name="conv"):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_conv_transpose = use_conv_transpose
+        self.name = name
+
+        self.conv = None
+        if use_conv_transpose:
+            self.conv = nn.ConvTranspose1d(channels, self.out_channels, 4, 2, 1)
+        elif use_conv:
+            self.conv = nn.Conv1d(self.channels, self.out_channels, 3, padding=1)
+
+    def forward(self, inputs):
+        assert inputs.shape[1] == self.channels
+        if self.use_conv_transpose:
+            return self.conv(inputs)
+
+        outputs = F.interpolate(inputs, scale_factor=2.0, mode="nearest")
+
+        if self.use_conv:
+            outputs = self.conv(outputs)
+
+        return outputs
+
+
+class ConditionalDecoder(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        channels=(256, 256),
+        dropout=0.0,
+        attention_head_dim=64,
+        n_blocks=4,
+        num_mid_blocks=12,
+        num_heads=8,
+        act_fn="gelu",
+    ):
+        """
+        This decoder requires an input with the same shape of the target. So, if your text content
+        is shorter or longer than the outputs, please re-sampling it before feeding to the decoder.
+        """
+        super().__init__()
+        channels = tuple(channels)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.time_embeddings = SinusoidalPosEmb(in_channels)
+        time_embed_dim = channels[0] * 4
+        self.time_mlp = TimestepEmbedding(
+            in_channels=in_channels,
+            time_embed_dim=time_embed_dim,
+            act_fn="silu",
+        )
+        self.down_blocks = nn.ModuleList([])
+        self.mid_blocks = nn.ModuleList([])
+        self.up_blocks = nn.ModuleList([])
+
+        output_channel = in_channels
+        for i in range(len(channels)):  # pylint: disable=consider-using-enumerate
+            input_channel = output_channel
+            output_channel = channels[i]
+            is_last = i == len(channels) - 1
+            resnet = ResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim)
+            transformer_blocks = nn.ModuleList(
+                [
+                    BasicTransformerBlock(
+                        dim=output_channel,
+                        num_attention_heads=num_heads,
+                        attention_head_dim=attention_head_dim,
+                        dropout=dropout,
+                        activation_fn=act_fn,
+                    )
+                    for _ in range(n_blocks)
+                ]
+            )
+            downsample = (
+                Downsample1D(output_channel) if not is_last else nn.Conv1d(output_channel, output_channel, 3, padding=1)
+            )
+            self.down_blocks.append(nn.ModuleList([resnet, transformer_blocks, downsample]))
+
+        for i in range(num_mid_blocks):
+            input_channel = channels[-1]
+            out_channels = channels[-1]
+            resnet = ResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim)
+
+            transformer_blocks = nn.ModuleList(
+                [
+                    BasicTransformerBlock(
+                        dim=output_channel,
+                        num_attention_heads=num_heads,
+                        attention_head_dim=attention_head_dim,
+                        dropout=dropout,
+                        activation_fn=act_fn,
+                    )
+                    for _ in range(n_blocks)
+                ]
+            )
+
+            self.mid_blocks.append(nn.ModuleList([resnet, transformer_blocks]))
+
+        channels = channels[::-1] + (channels[0],)
+        for i in range(len(channels) - 1):
+            input_channel = channels[i] * 2
+            output_channel = channels[i + 1]
+            is_last = i == len(channels) - 2
+            resnet = ResnetBlock1D(
+                dim=input_channel,
+                dim_out=output_channel,
+                time_emb_dim=time_embed_dim,
+            )
+            transformer_blocks = nn.ModuleList(
+                [
+                    BasicTransformerBlock(
+                        dim=output_channel,
+                        num_attention_heads=num_heads,
+                        attention_head_dim=attention_head_dim,
+                        dropout=dropout,
+                        activation_fn=act_fn,
+                    )
+                    for _ in range(n_blocks)
+                ]
+            )
+            upsample = (
+                Upsample1D(output_channel, use_conv_transpose=True)
+                if not is_last
+                else nn.Conv1d(output_channel, output_channel, 3, padding=1)
+            )
+            self.up_blocks.append(nn.ModuleList([resnet, transformer_blocks, upsample]))
+        self.final_block = Block1D(channels[-1], channels[-1])
+        self.final_proj = nn.Conv1d(channels[-1], self.out_channels, 1)
+        self.initialize_weights()
+
+
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv1d):
+                nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.GroupNorm):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def forward(self, x, mask, mu, t):
+        """Forward pass of the UNet1DConditional model.
+
+        Args:
+            x (torch.Tensor): shape (batch_size, in_channels, time)
+            mask (_type_): shape (batch_size, 1, time)
+            t (_type_): shape (batch_size)
+            spks (_type_, optional): shape: (batch_size, condition_channels). Defaults to None.
+            cond (_type_, optional): placeholder for future use. Defaults to None.
+
+        Raises:
+            ValueError: _description_
+            ValueError: _description_
+
+        Returns:
+            _type_: _description_
+        """
+
+        t = self.time_embeddings(t)
+        t = self.time_mlp(t)
+
+        x = pack([x, mu], "b * t")[0]
+
+        hiddens = []
+        masks = [mask]
+        for resnet, transformer_blocks, downsample in self.down_blocks:
+            mask_down = masks[-1]
+            x = resnet(x, mask_down, t)
+            x = rearrange(x, "b c t -> b t c").contiguous()
+            attn_mask = torch.matmul(mask_down.transpose(1, 2).contiguous(), mask_down)
+            for transformer_block in transformer_blocks:
+                x = transformer_block(
+                    hidden_states=x,
+                    attention_mask=attn_mask,
+                    timestep=t,
+                )
+            x = rearrange(x, "b t c -> b c t").contiguous()
+            hiddens.append(x)  # Save hidden states for skip connections
+            x = downsample(x * mask_down)
+            masks.append(mask_down[:, :, ::2])
+        masks = masks[:-1]
+        mask_mid = masks[-1]
+
+        for resnet, transformer_blocks in self.mid_blocks:
+            x = resnet(x, mask_mid, t)
+            x = rearrange(x, "b c t -> b t c").contiguous()
+            attn_mask = torch.matmul(mask_mid.transpose(1, 2).contiguous(), mask_mid)
+            for transformer_block in transformer_blocks:
+                x = transformer_block(
+                    hidden_states=x,
+                    attention_mask=attn_mask,
+                    timestep=t,
+                )
+            x = rearrange(x, "b t c -> b c t").contiguous()
+
+        for resnet, transformer_blocks, upsample in self.up_blocks:
+            mask_up = masks.pop()
+            skip = hiddens.pop()
+            x = pack([x[:, :, :skip.shape[-1]], skip], "b * t")[0]
+            x = resnet(x, mask_up, t)
+            x = rearrange(x, "b c t -> b t c").contiguous()
+            attn_mask = torch.matmul(mask_up.transpose(1, 2).contiguous(), mask_up)
+            for transformer_block in transformer_blocks:
+                x = transformer_block(
+                    hidden_states=x,
+                    attention_mask=attn_mask,
+                    timestep=t,
+                )
+            x = rearrange(x, "b t c -> b c t").contiguous()
+            x = upsample(x * mask_up)
+        x = self.final_block(x, mask_up)
+        output = self.final_proj(x * mask_up)
+        return output * mask
+
+
+class ConditionalCFM(nn.Module):
+    def __init__(self, 
+                 estimator: nn.Module,
+                 t_scheduler: str = "cosine",
+                 inference_cfg_rate: float = 0.7,
+                 ):
+        super().__init__()
+        self.estimator = estimator
+        self.t_scheduler = t_scheduler
+        self.inference_cfg_rate = inference_cfg_rate
+    
+    def solve_euler(self, x, t_span, mu, mask):
+        t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0]
+
+        # I am storing this because I can later plot it by putting a debugger here and saving it to a file
+        # Or in future might add like a return_all_steps flag
+        sol = []
+
+        for step in range(1, len(t_span)):
+            dphi_dt = self.estimator(x, mask, mu, t)
+            # Classifier-Free Guidance inference introduced in VoiceBox
+            if self.inference_cfg_rate > 0:
+                cfg_dphi_dt = self.estimator(x, mask, torch.zeros_like(mu), t)
+                dphi_dt = ((1.0 + self.inference_cfg_rate) * dphi_dt -
+                           self.inference_cfg_rate * cfg_dphi_dt)
+            x = x + dt * dphi_dt
+            t = t + dt
+            sol.append(x)
+            if step < len(t_span) - 1:
+                dt = t_span[step + 1] - t
+
+        return sol[-1]
+    
+    def inference(self, mu, mask, n_timesteps, temperature: float=1.0):
+        z = torch.randn_like(mu) * temperature
+        t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device)
+        if self.t_scheduler == 'cosine':
+            t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)
+        return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask)
+    

fireredtts/modules/flow/flow_model.py

@@ -0,0 +1,89 @@
+import time
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from fireredtts.modules.flow.utils import make_pad_mask
+
+
+class InterpolateRegulator(nn.Module):
+    def __init__(
+            self,
+            channels: int = 512,
+            num_blocks: int = 4,
+            groups: int = 1,
+    ):
+        super().__init__()
+        model = []
+        for _ in range(num_blocks):
+            model.extend([
+                nn.Conv1d(channels, channels, 3, 1, 1),
+                nn.GroupNorm(groups, channels),
+                nn.Mish(),
+            ])
+        model.append(
+            nn.Conv1d(channels, channels, 1, 1)
+        )
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x, ylens=None):
+        # x in (B, T, D)
+        mask = (~make_pad_mask(ylens)).to(x).unsqueeze(-1)
+        x = F.interpolate(x.transpose(1, 2).contiguous(), size=ylens.max(), mode='nearest')
+        out = self.model(x).transpose(1, 2).contiguous()
+        olens = ylens
+        return out * mask, olens
+
+
+class CrossAttnFlowMatching(nn.Module):
+    def __init__(self, 
+                 output_size: int,
+                 input_embedding: nn.Module,
+                 encoder: nn.Module,
+                 length_regulator: nn.Module,
+                 mel_encoder: nn.Module,
+                 decoder: nn.Module,
+                 ):
+        super().__init__()
+        self.input_embedding = input_embedding
+        self.encoder = encoder
+        self.length_regulator = length_regulator
+        self.encoder_proj = torch.nn.Linear(self.encoder.output_size, output_size)
+        self.prompt_prenet = mel_encoder
+        self.decoder = decoder
+
+    def inference(self, 
+                  token: torch.Tensor, 
+                  token_len: torch.Tensor, 
+                  prompt_mel: torch.Tensor, 
+                  prompt_mel_len: torch.Tensor, 
+                  n_timesteps:int=10,
+        ):
+        # prompt projection
+        prompt_feat = self.prompt_prenet(prompt_mel)
+        prompt_feat_len = torch.ceil(prompt_mel_len/self.prompt_prenet.reduction_rate).long()
+
+        # concat text and prompt_text
+        mask = (~make_pad_mask(token_len)).float().unsqueeze(-1).to(token_len.device)
+        token = self.input_embedding(torch.clamp(token, min=0)) * mask
+
+        # 40ms shift to 10ms shift
+        feat_len = (token_len *4).int()
+
+        # first encoder
+        h, _ = self.encoder(token, token_len, prompt_feat, prompt_feat_len)
+        # length regulate
+        h, _ = self.length_regulator(h, feat_len)
+        # final projection
+        h = self.encoder_proj(h)
+        
+        mask = (~make_pad_mask(feat_len)).to(h)
+
+        feat = self.decoder.inference(
+            mu=h.transpose(1, 2).contiguous(),
+            mask=mask.unsqueeze(1),
+            n_timesteps=n_timesteps,
+        )
+        return feat
+
+

fireredtts/modules/flow/mel_encoder.py

@@ -0,0 +1,170 @@
+import typing as tp
+import torch
+import torch.nn as nn
+
+
+class ConvLayer(nn.Module):
+    def __init__(self, 
+                 in_channels:int,
+                 out_channels:int,
+                 kernel_size:int,
+                 stride:int,
+                 activation:str="GELU",
+                 dropout_rate:float=0.0,
+                 ):
+        super().__init__()
+        self.conv = nn.Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=(kernel_size-stride)//2,
+        )
+        self.drop = nn.Dropout(dropout_rate)
+        self.norm = nn.LayerNorm(out_channels)
+        self.activ = getattr(nn, activation)()
+
+    def forward(self, x:torch.Tensor):
+        """
+        Args:
+            x: (b, t, c)
+        Return:
+            x: (b, t, c)
+        """
+        x = x.transpose(2, 1)
+        x = self.conv(x)
+        x = x.transpose(2, 1)
+        x = self.drop(x)
+        x = self.norm(x)
+        x = self.activ(x)
+        return x
+
+
+class ResidualConvLayer(nn.Module):
+    def __init__(self,
+                 hidden_channels:int,
+                 n_layers:int=2,
+                 kernel_size:int=5,
+                 activation:str="GELU",
+                 dropout_rate:float=0.0,
+                 ):
+        super().__init__()
+        layers = [
+            ConvLayer(hidden_channels, hidden_channels, kernel_size, 1, activation, dropout_rate)
+            for _ in range(n_layers)
+        ]
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x:torch.Tensor):
+        """
+        Args:
+            x: (b, t, c)
+        Returns:
+            x: (b, t, c)
+        """
+        return x + self.layers(x)
+
+
+class ResidualConvBlock(nn.Module):
+    def __init__(self, 
+                 in_channels:int,
+                 hidden_channels:int,
+                 out_channels:int,
+                 n_layers:int=2,
+                 n_blocks:int=5,
+                 middle_layer:tp.Optional[nn.Module]=None,
+                 kernel_size:int=5,
+                 activation:str="GELU",
+                 dropout_rate:float=0.0,
+                 ):
+        super().__init__()
+        self.in_proj = nn.Conv1d(
+            in_channels,
+            hidden_channels,
+            kernel_size=kernel_size,
+            stride=1,
+            padding=(kernel_size-1)//2,
+        ) if in_channels != hidden_channels else nn.Identity()
+
+        self.conv1 = nn.Sequential(*[
+            ResidualConvLayer(hidden_channels, n_layers, kernel_size, activation, dropout_rate)
+            for _ in range(n_blocks)
+        ])
+
+        if middle_layer is None:
+            self.middle_layer = nn.Identity()
+        elif isinstance(middle_layer, nn.Module):
+            self.middle_layer = middle_layer
+        else:
+            raise TypeError("unknown middle layer type:{}".format(type(middle_layer)))
+        
+        self.conv2 = nn.Sequential(*[
+            ResidualConvLayer(hidden_channels, n_layers, kernel_size, activation, dropout_rate)
+            for _ in range(n_blocks)
+        ])
+
+        self.out_proj = nn.Conv1d(
+            hidden_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=1,
+            padding=(kernel_size-1)//2,
+        ) if out_channels != hidden_channels else nn.Identity()
+
+    def forward(self, x:torch.Tensor, **middle_layer_kwargs):
+        """
+        Args:
+            x: (b, t1, c)
+        Return:
+            x: (b, t2, c)
+        """
+        x = self.in_proj(x.transpose(2, 1)).transpose(2, 1)
+        x = self.conv1(x)
+        if isinstance(self.middle_layer, nn.MaxPool1d) or isinstance(self.middle_layer, nn.Conv1d):
+            x = self.middle_layer(x.transpose(2, 1)).transpose(2, 1)
+        elif isinstance(self.middle_layer, nn.Identity):
+            x = self.middle_layer(x)
+        else:
+            # incase of phoneme-pooling layer
+            x = self.middle_layer(x, **middle_layer_kwargs)
+        x = self.conv2(x)
+        x = self.out_proj(x.transpose(2, 1)).transpose(2, 1)
+        return x
+
+
+class MelReduceEncoder(nn.Module):
+    def __init__(self, 
+                 in_channels:int,
+                 out_channels:int,
+                 hidden_channels:int=384,
+                 reduction_rate:int=4,
+                 n_layers:int=2,
+                 n_blocks:int=5,
+                 kernel_size:int=3,
+                 activation:str="GELU",
+                 dropout:float=0.0,
+                 ):
+        super().__init__()
+        self.reduction_rate = reduction_rate
+        middle_conv = nn.Conv1d(
+            in_channels=hidden_channels,
+            out_channels=hidden_channels,
+            kernel_size=reduction_rate,
+            stride=reduction_rate,
+            padding=0
+        )
+        self.encoder = ResidualConvBlock(
+            in_channels=in_channels,
+            hidden_channels=hidden_channels,
+            out_channels=out_channels,
+            n_layers=n_layers,
+            n_blocks=n_blocks,
+            middle_layer=middle_conv,
+            kernel_size=kernel_size,
+            activation=activation,
+            dropout_rate=dropout
+        )
+    
+    def forward(self, x:torch.Tensor):
+        return self.encoder(x)
+

fireredtts/modules/flow/mel_spectrogram.py

@@ -0,0 +1,132 @@
+from functools import partial
+import torch
+import numpy as np
+import librosa
+from librosa.filters import mel as librosa_mel_fn
+from torchaudio.functional import resample as ta_resample_fn
+
+MAX_WAV_VALUE = 32767.0  # NOTE: 32768.0 -1 to prevent int16 overflow (results in popping sound in corner cases)
+
+
+def dynamic_range_compression(x, C=1, clip_val=1e-5):
+    return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)
+
+
+def dynamic_range_decompression(x, C=1):
+    return np.exp(x) / C
+
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def dynamic_range_decompression_torch(x, C=1):
+    return torch.exp(x) / C
+
+
+def spectral_normalize_torch(magnitudes):
+    output = dynamic_range_compression_torch(magnitudes)
+    return output
+
+
+def spectral_de_normalize_torch(magnitudes):
+    output = dynamic_range_decompression_torch(magnitudes)
+    return output
+
+
+mel_basis = {}
+hann_window = {}
+
+
+def mel_spectrogram(
+    y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False
+):
+    global mel_basis, hann_window
+    if fmax not in mel_basis:
+        mel = librosa_mel_fn(
+            sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax
+        )
+        str_key_mel_basis = str(fmax) + "_" + str(y.device)
+        mel_basis[str_key_mel_basis] = torch.from_numpy(mel).float().to(y.device)
+        hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device)
+
+    y = torch.nn.functional.pad(
+        y.unsqueeze(1),
+        (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),
+        mode="reflect",
+    )
+    y = y.squeeze(1)
+
+    # complex tensor as default, then use view_as_real for future pytorch compatibility
+    spec = torch.stft(
+        y,
+        n_fft,
+        hop_length=hop_size,
+        win_length=win_size,
+        window=hann_window[str(y.device)],
+        center=center,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=True,
+    )
+    spec = torch.view_as_real(spec)
+    spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9))
+
+    spec = torch.matmul(mel_basis[str_key_mel_basis], spec)
+    spec = spectral_normalize_torch(spec)
+
+    return spec
+
+
+kaiser_best_resampling_fn = partial(
+    ta_resample_fn, 
+    resampling_method="sinc_interp_kaiser", # DO NOT CHANGE!
+    rolloff=0.917347, # DO NOT CHANGE!
+    beta=12.9846, # DO NOT CHANGE!
+    lowpass_filter_width=50, # DO NOT CHANGE!
+)
+
+
+class MelSpectrogramExtractor(object):
+    def __init__(
+        self,
+        n_fft=1024,
+        win_size=1024,
+        num_mels=100,
+        hop_size=160,
+        sampling_rate=16000,
+        fmin=0,
+        fmax=None,
+    ):
+        self.n_fft = n_fft
+        self.win_size = win_size
+        self.num_mels = num_mels
+        self.hop_size = hop_size
+        self.sampling_rate = sampling_rate
+        self.fmin = fmin
+        self.fmax = fmax
+
+    def __call__(self, wav_path) -> np.ndarray:
+        wav_data, wav_sr = librosa.load(wav_path, sr=None, mono=True)
+        wav_data = torch.from_numpy(wav_data.copy()).unsqueeze(0)
+        # for 16k wavs, up-downsample to reduce artifects
+        if wav_sr == self.sampling_rate:
+            wav_data = kaiser_best_resampling_fn(wav_data, orig_freq=wav_sr, new_freq=24000)
+            wav_data = kaiser_best_resampling_fn(wav_data, orig_freq=24000, new_freq=self.sampling_rate)
+        else:
+            wav_data = kaiser_best_resampling_fn(wav_data, orig_freq=wav_sr, new_freq=self.sampling_rate)        
+
+        # (1, num_mels, t)
+        mel = mel_spectrogram(
+            wav_data,
+            self.n_fft,
+            self.num_mels,
+            self.sampling_rate,
+            self.hop_size,
+            self.win_size,
+            self.fmin,
+            self.fmax,
+        )
+        mel = mel.squeeze(0).transpose(1, 0)
+        return mel  # (t, num_mels)

fireredtts/modules/flow/transformer.py

@@ -0,0 +1,249 @@
+from typing import Any, Dict, Optional
+
+import torch
+import torch.nn as nn
+from diffusers.models.attention import (
+    GEGLU,
+    GELU,
+    AdaLayerNorm,
+    AdaLayerNormZero,
+    ApproximateGELU,
+)
+from diffusers.models.attention_processor import Attention
+# from diffusers.models.lora import LoRACompatibleLinear
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+
+
+class FeedForward(nn.Module):
+    r"""
+    A feed-forward layer.
+
+    Parameters:
+        dim (`int`): The number of channels in the input.
+        dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.
+        mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        final_dropout (`bool` *optional*, defaults to False): Apply a final dropout.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        dim_out: Optional[int] = None,
+        mult: int = 4,
+        dropout: float = 0.0,
+        activation_fn: str = "geglu",
+        final_dropout: bool = False,
+    ):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+
+        if activation_fn == "gelu":
+            act_fn = GELU(dim, inner_dim)
+        if activation_fn == "gelu-approximate":
+            act_fn = GELU(dim, inner_dim, approximate="tanh")
+        elif activation_fn == "geglu":
+            act_fn = GEGLU(dim, inner_dim)
+        elif activation_fn == "geglu-approximate":
+            act_fn = ApproximateGELU(dim, inner_dim)
+
+        self.net = nn.ModuleList([])
+        # project in
+        self.net.append(act_fn)
+        # project dropout
+        self.net.append(nn.Dropout(dropout))
+        # project out
+        self.net.append(nn.Linear(inner_dim, dim_out))
+        # self.net.append(LoRACompatibleLinear(inner_dim, dim_out))
+        # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
+        if final_dropout:
+            self.net.append(nn.Dropout(dropout))
+
+    def forward(self, hidden_states):
+        for module in self.net:
+            hidden_states = module(hidden_states)
+        return hidden_states
+
+
+@maybe_allow_in_graph
+class BasicTransformerBlock(nn.Module):
+    r"""
+    A basic Transformer block.
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
+        only_cross_attention (`bool`, *optional*):
+            Whether to use only cross-attention layers. In this case two cross attention layers are used.
+        double_self_attention (`bool`, *optional*):
+            Whether to use two self-attention layers. In this case no cross attention layers are used.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        num_embeds_ada_norm (:
+            obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
+        attention_bias (:
+            obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        dropout=0.0,
+        cross_attention_dim: Optional[int] = None,
+        activation_fn: str = "geglu",
+        num_embeds_ada_norm: Optional[int] = None,
+        attention_bias: bool = False,
+        only_cross_attention: bool = False,
+        double_self_attention: bool = False,
+        upcast_attention: bool = False,
+        norm_elementwise_affine: bool = True,
+        norm_type: str = "layer_norm",
+        final_dropout: bool = False,
+    ):
+        super().__init__()
+        self.only_cross_attention = only_cross_attention
+
+        self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
+        self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm"
+
+        if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
+            raise ValueError(
+                f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"
+                f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."
+            )
+
+        # Define 3 blocks. Each block has its own normalization layer.
+        # 1. Self-Attn
+        if self.use_ada_layer_norm:
+            self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm)
+        elif self.use_ada_layer_norm_zero:
+            self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)
+        else:
+            self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
+        self.attn1 = Attention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            dropout=dropout,
+            bias=attention_bias,
+            cross_attention_dim=cross_attention_dim if only_cross_attention else None,
+            upcast_attention=upcast_attention,
+        )
+
+        # 2. Cross-Attn
+        if cross_attention_dim is not None or double_self_attention:
+            # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
+            # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
+            # the second cross attention block.
+            self.norm2 = (
+                AdaLayerNorm(dim, num_embeds_ada_norm)
+                if self.use_ada_layer_norm
+                else nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
+            )
+            self.attn2 = Attention(
+                query_dim=dim,
+                cross_attention_dim=cross_attention_dim if not double_self_attention else None,
+                heads=num_attention_heads,
+                dim_head=attention_head_dim,
+                dropout=dropout,
+                bias=attention_bias,
+                upcast_attention=upcast_attention,
+                # scale_qk=False, # uncomment this to not to use flash attention
+            )  # is self-attn if encoder_hidden_states is none
+        else:
+            self.norm2 = None
+            self.attn2 = None
+
+        # 3. Feed-forward
+        self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
+        self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout)
+
+        # let chunk size default to None
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int):
+        # Sets chunk feed-forward
+        self._chunk_size = chunk_size
+        self._chunk_dim = dim
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        timestep: Optional[torch.LongTensor] = None,
+        cross_attention_kwargs: Dict[str, Any] = None,
+        class_labels: Optional[torch.LongTensor] = None,
+    ):
+        # Notice that normalization is always applied before the real computation in the following blocks.
+        # 1. Self-Attention
+        if self.use_ada_layer_norm:
+            norm_hidden_states = self.norm1(hidden_states, timestep)
+        elif self.use_ada_layer_norm_zero:
+            norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+                hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
+            )
+        else:
+            norm_hidden_states = self.norm1(hidden_states)
+
+        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
+
+        attn_output = self.attn1(
+            norm_hidden_states,
+            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
+            attention_mask=encoder_attention_mask if self.only_cross_attention else attention_mask,
+            **cross_attention_kwargs,
+        )
+        if self.use_ada_layer_norm_zero:
+            attn_output = gate_msa.unsqueeze(1) * attn_output
+        hidden_states = attn_output + hidden_states
+
+        # 2. Cross-Attention
+        if self.attn2 is not None:
+            norm_hidden_states = (
+                self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)
+            )
+
+            attn_output = self.attn2(
+                norm_hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=encoder_attention_mask,
+                **cross_attention_kwargs,
+            )
+            hidden_states = attn_output + hidden_states
+
+        # 3. Feed-forward
+        norm_hidden_states = self.norm3(hidden_states)
+
+        if self.use_ada_layer_norm_zero:
+            norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+        if self._chunk_size is not None:
+            # "feed_forward_chunk_size" can be used to save memory
+            if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
+                raise ValueError(
+                    f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
+                )
+
+            num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
+            ff_output = torch.cat(
+                [self.ff(hid_slice) for hid_slice in norm_hidden_states.chunk(num_chunks, dim=self._chunk_dim)],
+                dim=self._chunk_dim,
+            )
+        else:
+            ff_output = self.ff(norm_hidden_states)
+
+        if self.use_ada_layer_norm_zero:
+            ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+        hidden_states = ff_output + hidden_states
+
+        return hidden_states

fireredtts/modules/flow/utils.py

@@ -0,0 +1,30 @@
+import torch
+
+
+def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
+    """Make mask tensor containing indices of padded part.
+
+    See description of make_non_pad_mask.
+
+    Args:
+        lengths (torch.Tensor): Batch of lengths (B,).
+    Returns:
+        torch.Tensor: Mask tensor containing indices of padded part.
+
+    Examples:
+        >>> lengths = [5, 3, 2]
+        >>> make_pad_mask(lengths)
+        masks = [[0, 0, 0, 0 ,0],
+                 [0, 0, 0, 1, 1],
+                 [0, 0, 1, 1, 1]]
+    """
+    batch_size = lengths.size(0)
+    max_len = max_len if max_len > 0 else lengths.max().item()
+    seq_range = torch.arange(0,
+                             max_len,
+                             dtype=torch.int64,
+                             device=lengths.device)
+    seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
+    seq_length_expand = lengths.unsqueeze(-1)
+    mask = seq_range_expand >= seq_length_expand
+    return mask

fireredtts/modules/gpt/gpt.py

@@ -0,0 +1,356 @@
+# ported from: https://github.com/neonbjb/tortoise-tts
+# ported from: https://github.com/coqui-ai/TTS/blob/dev/TTS/tts/layers/xtts/gpt.py
+
+import functools
+import math
+import random
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import GPT2Config, GPT2Model, 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
+        )
+
+        # 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
+        )
+
+
+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):
+        super().__init__()
+        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)
+
+    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.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.
+    """
+    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)
+    )
+    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,
+        checkpointing=False,
+        label_smoothing=0.0,
+    ):
+        """
+        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_dropout = nn.Dropout1d(0.1)
+        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,
+        )
+
+        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)
+
+        # reference_embedding
+        self.reference_embedding = nn.Sequential(
+            nn.Linear(512, 256),
+            nn.Tanh(),
+            nn.Linear(256, self.model_dim),
+        )
+
+    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
+
+    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] :]

fireredtts/modules/text_normalizer/normalize.py

@@ -0,0 +1,178 @@
+import re
+import regex
+import inflect
+import unicodedata
+from lingua import Language, LanguageDetectorBuilder
+from builtins import str as unicode
+
+from tn.chinese.normalizer import Normalizer as ZhNormalizer
+from tn.english.normalizer import Normalizer as EnNormalizer
+
+from fireredtts.modules.text_normalizer.regex_common import *
+from fireredtts.modules.text_normalizer.utils import *
+
+
+def preprocess_text(sentence):
+    # preprocessing
+    sentence = bytes(sentence, "utf-8").decode("utf-8", "ignore")
+    sentence = regex.sub("[\p{Cf}--[\u200d]]", "", sentence, flags=regex.V1)
+    sentence = regex.sub("\p{Co}", "", sentence)
+    sentence = sentence.replace("\u00a0", " ")
+    sentence = sentence.replace("\ufffd", "")
+    sentence = regex.sub("\p{Zl}", "\n", sentence)
+    sentence = regex.sub("\p{Zp}", "\n", sentence)
+
+    sentence = unicode(sentence)
+    sentence = "".join(
+        char
+        for char in unicodedata.normalize("NFD", sentence)
+        if unicodedata.category(char) != "Mn"
+    )  # Strip accents
+
+    sentence = strip_kaomoji(sentence)
+    # full to half with exemption (to be converted after number TN): 。，：
+    sentence = f2b(sentence, exemption="。，：")
+
+    # clean spaces
+    sentence = sentence.replace("\n", "，")
+    sentence = sentence.replace("\t", "，")
+    sentence = sentence.replace("\r", "，")
+    sentence = re.sub(r"[。.]{3,}", "…", sentence)
+    sentence = re.sub(r"[…⋯]{1,}", "…", sentence)
+    sentence = re.sub(r"[ ]+", " ", sentence)
+    sentence = sentence.strip()
+
+    # punctuation reduction
+    result = ""
+    for idx, char in enumerate(sentence):
+        if char in symbol_reduction:
+            char = symbol_reduction[char]
+
+        if char == " ":
+            if idx == 0:
+                continue
+            if is_chinese(sentence[idx + 1]) and (
+                is_chinese(sentence[idx - 1]) or sentence[idx - 1] in '") '
+            ):
+                result += "，"
+            else:
+                result += " "
+            continue
+
+        if is_valid_char(char):
+            result += char
+    result = re.sub(r"[ ]+", " ", result)
+    return result
+
+
+def rettt(sentence):
+    # handle abbreviations for all languages
+    sentence = sentence.replace("&nd", "and")
+    sentence = sentence.replace("Jan.", "january")
+    sentence = sentence.replace("Feb.", "febrary")
+    sentence = sentence.replace("Mar.", "march")
+    sentence = sentence.replace("Apr.", "april")
+    sentence = sentence.replace("May.", "may")
+    sentence = sentence.replace("Jun.", "june")
+    sentence = sentence.replace("Jul.", "july")
+    sentence = sentence.replace("Aug.", "august")
+    sentence = sentence.replace("Sept.", "september")
+    sentence = sentence.replace("Sep.", "september")
+    sentence = sentence.replace("Oct.", "october")
+    sentence = sentence.replace("Nov.", "november")
+    sentence = sentence.replace("Dec.", "december")
+    sentence = sentence.replace("Mon.", "monday")
+    sentence = sentence.replace("Tues.", "tuesday")
+    sentence = sentence.replace("Wed.", "wednesday")
+    sentence = sentence.replace("Thur.", "thursday")
+    sentence = sentence.replace("Fri.", "friday")
+    sentence = sentence.replace("Sat.", "saturday")
+    if sentence != "Sun.":
+        sentence = sentence.replace("Sun.", "sunday")
+    sentence = re.sub(r" St\. ([A-Z])", r" saint \1", sentence)
+    sentence = re.sub(r" St\.", " street", sentence)
+    sentence = re.sub(r" Rd\.", " road", sentence)
+    sentence = re.sub(r"[Aa]\.[Mm]\.", "A_M", sentence)
+    sentence = re.sub(r"[Pp]\.[Mm]\.", "P_M", sentence)
+    sentence = re.sub(r"[Bb]\.[Cc]\.", "B_C", sentence)
+    sentence = re.sub(r"[Ad]\.[Dd]\.", "A_D", sentence)
+    sentence = sentence.replace("Mr.", "mister")
+    sentence = sentence.replace("Ms.", "miss")
+    sentence = sentence.replace("Mrs.", "misses")
+    sentence = sentence.replace("Ph.D", "P_H_D")
+    sentence = sentence.replace("i.e.", "that is")
+    sentence = sentence.replace("e.g.", "for example")
+    sentence = sentence.replace("btw.", "by the way")
+    sentence = sentence.replace("btw", "by the way")
+    sentence = sentence.replace("b.t.w.", "by the way")
+    sentence = sentence.replace("@", " at ")
+    return sentence
+
+
+class TextNormalizer:
+    def __init__(self):
+        self.language_detector = LanguageDetectorBuilder.from_languages(
+            Language.ENGLISH, Language.CHINESE
+        ).build()
+        self.zh_normalizer = ZhNormalizer()
+        self.en_normalizer = EnNormalizer()
+        self.inflect_parser = inflect.engine()
+        self.lang2token = {Language.ENGLISH: "en", Language.CHINESE: "zh"}
+
+    def tn(self, text):
+        text = preprocess_text(text)
+        text = rettt(text)  # regex replacements
+        # for non chinese languages
+        language = self.language_detector.detect_language_of(text)
+        # enforce chinese if text contains any chinese character
+        if contains_chinese(text):
+            language = Language.CHINESE
+        text_lang = self.lang2token.get(language, "zh")
+
+        if is_upper_eng_and_digit(text):
+            language = Language.CHINESE
+
+        if language == Language.CHINESE:
+            text = self.zh_normalizer.normalize(text)
+            text = text.replace("\n", "")
+            text = re.sub(r"[，,]+$", "。", text)
+        else:
+            text = re.sub(r"[^ 0-9A-Za-z\[\]'.,:?!_\-]", "", text)
+            text = self.en_normalizer.normalize(text)
+            # fallback number normalization
+            pieces = re.split(r"(\d+)", text)
+            text = "".join(
+                [
+                    self.inflect_parser.number_to_words(p) if p.isnumeric() else p
+                    for p in pieces
+                    if len(p) > 0
+                ]
+            )
+
+        # cleanup
+        text = text.replace("_", " ")
+        text = re.sub(r"[ ]+", " ", text)
+
+        # spell caplital words
+        pieces = re.split(r"([A-Z]{2,4}|[ ])", text)
+        for idx, p in enumerate(pieces):
+            if re.match("[A-Z]{2,4}", p):
+                pieces[idx] = " ".join(p)
+        text = " ".join([p for p in pieces if p != " "])
+
+        # post TN full to half
+        text = text.replace("。", ".")
+        text = text.replace("，", ",")
+        text = text.replace("：", ":")
+
+        # model limitations
+        text = text.lower().strip()
+        text = text.replace('"', "")
+        text = text.replace("·", " ")
+        text = re.sub("[…~、！，？：；!?:;]+", ",", text)
+        text = re.sub("[,]+", ",", text)
+        text = re.sub(r"[,. ]+$", ".", text)
+        if len(text) > 0 and text[-1] != ".":
+            text = text + "."
+
+        return text, text_lang

fireredtts/modules/text_normalizer/regex_common.py

@@ -0,0 +1,23 @@
+import re
+
+kaomoji_regex = re.compile(
+    r"[oヽwΣ┗╰O︿Ψ凸]?[(|≡*（].{0,4}[Д✿_▽→≧﹏`∩⊙∇☆≡๑〃′エ≦▔＠﹁εヘ•́ω益‿≖ฺ皿•̀艹￣△|ﾟ].{0,5}[|≡*)）][┛ブ凸cｄd︴oOΨ︿w╯ノ]?"
+)
+chinese_regex = re.compile(r"[\u4e00-\u9fa5]")
+digit_regex = re.compile(r"(\\d+)(\\.\\d+)?", re.UNICODE)
+
+chinese_char_regex = re.compile(r"^[\u4e00-\u9fa5]$", re.UNICODE)
+eng_and_digit_char_regex = re.compile(r"^[0-9.,A-Za-z]+$", re.UNICODE)
+upper_eng_and_digit_regex = re.compile(r"^[ 0-9A-Z\"'.,:?!\-]+$", re.UNICODE)
+valid_char_regex = re.compile(
+    r"[\t\r\n ]|"
+    r"[\u4e00-\u9fa5]|"
+    r"\u0080|[\u20a0-\u20bf]|\u00a2|\u00a3|\u00a5|\uffe0|\uffe1|\uffe5|\uffe6|"
+    r"\u3000|\u3002|\u00b7|\u2014|\u2019|\u2026|\uff01|\uff1f|\uff0e|\uff1a|\uff1b|\uff0b|\uff0c|\uff0d|\uff0f|[\ufe10-\ufe16]|[\ufe50-\ufe51]|[\ufe55-\ufe57]|\ufe6a|"
+    r"[\u0030-\u0039]|"
+    r"[\u0391-\u03c9]|"
+    r"[\u00b0-\u00b3]|[\u2015-\u2018]|[\u3000-\u303f]|"
+    r"[\u0022-\u002f\u003a-\u003e\u0040\u005b-\u0060\u007b-\u007e]|"
+    r"[\uff21-\uff3a]|[\uff41-\uff5a]|[\u0041-\u005a]|[\u0061-\u007a]",
+    re.UNICODE,
+)

fireredtts/modules/text_normalizer/utils.py

@@ -0,0 +1,121 @@
+from fireredtts.modules.text_normalizer.regex_common import *
+
+
+def contains_chinese(text):
+    return bool(chinese_regex.search(text))
+
+
+def strip_kaomoji(text):
+    return kaomoji_regex.sub(" ", text)
+
+
+def is_chinese(char):
+    return chinese_char_regex.match(char)
+
+
+def is_eng_and_digit(char):
+    return eng_and_digit_char_regex.match(char)
+
+
+def is_upper_eng_and_digit(text):
+    return upper_eng_and_digit_regex.match(text)
+
+
+def is_valid_char(char):
+    return valid_char_regex.match(char)
+
+
+def is_digit(text):
+    return digit_regex.match(text)
+
+
+def contains_chinese(text):
+    return bool(chinese_regex.search(text))
+
+
+def f2b(ustr, exemption="。，："):
+    half = []
+    for u in ustr:
+        num = ord(u)
+        if num == 0x3000:
+            half.append(" ")
+        elif u in exemption:  # exemption
+            half.append(u)
+        elif 0xFF01 <= num <= 0xFF5E:
+            num -= 0xFEE0
+            half.append(chr(num))
+        else:
+            half.append(u)
+    return "".join(half)
+
+
+symbol_reduction = {
+    "「": '"',
+    "」": '"',
+    "｀": '"',
+    "〝": '"',
+    "〞": '"',
+    "‟": '"',
+    "„": '"',
+    "｛": "(",
+    "｝": ")",
+    "【": "(",
+    "】": ")",
+    "〖": "(",
+    "〗": ")",
+    "〔": "(",
+    "〕": ")",
+    "〘": "(",
+    "〙": ")",
+    "《": "(",
+    "》": ")",
+    "｟": "(",
+    "｠": ")",
+    "〚": "(",
+    "〛": ")",
+    "『": '"',
+    "』": '"',
+    "｢": '"',
+    "｣": '"',
+    "{": "(",
+    "}": ")",
+    "〈": "(",
+    "〉": ")",
+    "•": "·",
+    "‧": "·",
+    "〰": "…",
+    "﹏": "…",
+    "〜": "~",
+    "～": "~",
+    "＋": "+",
+    "､": "、",
+    "｡": "。",
+    "︐": "，",
+    "﹐": "，",
+    "︑": "、",
+    "﹑": "、",
+    "︒": "。",
+    "︓": "：",
+    "﹕": "：",
+    "︔": "；",
+    "﹔": "；",
+    "︕": "！",
+    "﹗": "！",
+    "︖": "？",
+    "﹖": "？",
+    "﹙": "(",
+    "﹚": ")",
+    "﹪": "%",
+    "﹠": "&",
+    "＞": ">",
+    "|": "、",
+    "＝": "=",
+    "‐": "-",
+    "‑": "-",
+    "‒": "-",
+    "–": "-",
+    "—": "-",
+    "―": "-",
+    "％": "%",
+    "μ": "u",
+}

fireredtts/modules/tokenizer/tokenizer.py

@@ -0,0 +1,46 @@
+import os
+import torch
+from fireredtts.modules.tokenizer.whisper_tokenizer import get_tokenizer
+from fireredtts.modules.text_normalizer.normalize import TextNormalizer
+
+
+DEFAULT_VOCAB_FILE = os.path.join(
+    os.path.dirname(os.path.realpath(__file__)), "../data/tokenizer.json"
+)
+
+
+class VoiceBpeTokenizer:
+    def __init__(self):
+        self.tokenizer = get_tokenizer(multilingual=True)
+        self.tn_engine = TextNormalizer()
+
+    def redtts_text_cleaner(self, text):
+        text = text.strip()
+        text, text_lang = self.tn_engine.tn(text)
+        # print("---text after tn:", text)
+        return text, text_lang
+
+    def encode(self, text, lang="auto"):
+        text, text_lang = self.redtts_text_cleaner(text=text)
+        if lang == "auto":
+            lang = text_lang
+        text = f"[{lang}]{text}"
+        return self.tokenizer.encode(text)
+
+    def decode(self, seq):
+        if isinstance(seq, torch.Tensor):
+            seq = seq.cpu().numpy()
+        text = self.tokenizer.decode(seq)
+        return text
+
+    def __len__(self):
+        return self.tokenizer.get_vocab_size()
+
+    def get_number_tokens(self):
+        return self.tokenizer.get_vocab_size()
+
+
+if __name__ == "__main__":
+    tok = VoiceBpeTokenizer()
+    codes = tok.encode("我、真是hello USA啊？谢谢你world！")
+    print([tok.decode([c]) for c in codes])

fireredtts/modules/tokenizer/whisper_tokenizer.py

@@ -0,0 +1,456 @@
+# adapted from https://github.com/openai/whisper/blob/main/whisper/tokenizer.py
+# Copyright (c) 2022 OpenAI
+
+# MIT License for this file
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import base64
+import os
+import string
+from dataclasses import dataclass, field
+from functools import cached_property, lru_cache
+from typing import Dict, List, Optional, Tuple
+
+import tiktoken
+
+LANGUAGES = {
+    "en": "english",
+    "zh": "chinese",
+    "de": "german",
+    "es": "spanish",
+    "ru": "russian",
+    "ko": "korean",
+    "fr": "french",
+    "ja": "japanese",
+    "pt": "portuguese",
+    "tr": "turkish",
+    "pl": "polish",
+    "ca": "catalan",
+    "nl": "dutch",
+    "ar": "arabic",
+    "sv": "swedish",
+    "it": "italian",
+    "id": "indonesian",
+    "hi": "hindi",
+    "fi": "finnish",
+    "vi": "vietnamese",
+    "he": "hebrew",
+    "uk": "ukrainian",
+    "el": "greek",
+    "ms": "malay",
+    "cs": "czech",
+    "ro": "romanian",
+    "da": "danish",
+    "hu": "hungarian",
+    "ta": "tamil",
+    "no": "norwegian",
+    "th": "thai",
+    "ur": "urdu",
+    "hr": "croatian",
+    "bg": "bulgarian",
+    "lt": "lithuanian",
+    "la": "latin",
+    "mi": "maori",
+    "ml": "malayalam",
+    "cy": "welsh",
+    "sk": "slovak",
+    "te": "telugu",
+    "fa": "persian",
+    "lv": "latvian",
+    "bn": "bengali",
+    "sr": "serbian",
+    "az": "azerbaijani",
+    "sl": "slovenian",
+    "kn": "kannada",
+    "et": "estonian",
+    "mk": "macedonian",
+    "br": "breton",
+    "eu": "basque",
+    "is": "icelandic",
+    "hy": "armenian",
+    "ne": "nepali",
+    "mn": "mongolian",
+    "bs": "bosnian",
+    "kk": "kazakh",
+    "sq": "albanian",
+    "sw": "swahili",
+    "gl": "galician",
+    "mr": "marathi",
+    "pa": "punjabi",
+    "si": "sinhala",
+    "km": "khmer",
+    "sn": "shona",
+    "yo": "yoruba",
+    "so": "somali",
+    "af": "afrikaans",
+    "oc": "occitan",
+    "ka": "georgian",
+    "be": "belarusian",
+    "tg": "tajik",
+    "sd": "sindhi",
+    "gu": "gujarati",
+    "am": "amharic",
+    "yi": "yiddish",
+    "lo": "lao",
+    "uz": "uzbek",
+    "fo": "faroese",
+    "ht": "haitian creole",
+    "ps": "pashto",
+    "tk": "turkmen",
+    "nn": "nynorsk",
+    "mt": "maltese",
+    "sa": "sanskrit",
+    "lb": "luxembourgish",
+    "my": "myanmar",
+    "bo": "tibetan",
+    "tl": "tagalog",
+    "mg": "malagasy",
+    "as": "assamese",
+    "tt": "tatar",
+    "haw": "hawaiian",
+    "ln": "lingala",
+    "ha": "hausa",
+    "ba": "bashkir",
+    "jw": "javanese",
+    "su": "sundanese",
+    "yue": "cantonese",
+}
+
+# language code lookup by name, with a few language aliases
+TO_LANGUAGE_CODE = {
+    **{language: code for code, language in LANGUAGES.items()},
+    "burmese": "my",
+    "valencian": "ca",
+    "flemish": "nl",
+    "haitian": "ht",
+    "letzeburgesch": "lb",
+    "pushto": "ps",
+    "panjabi": "pa",
+    "moldavian": "ro",
+    "moldovan": "ro",
+    "sinhalese": "si",
+    "castilian": "es",
+    "mandarin": "zh",
+}
+
+
+@dataclass
+class Tokenizer:
+    """A thin wrapper around `tiktoken` providing quick access to special tokens"""
+
+    encoding: tiktoken.Encoding
+    num_languages: int
+    language: Optional[str] = None
+    task: Optional[str] = None
+    sot_sequence: Tuple[int] = ()
+    special_tokens: Dict[str, int] = field(default_factory=dict)
+
+    def __post_init__(self):
+        for special in self.encoding.special_tokens_set:
+            special_token = self.encoding.encode_single_token(special)
+            self.special_tokens[special] = special_token
+
+        sot: int = self.special_tokens["[startoftranscript]"]
+        translate: int = self.special_tokens["[translate]"]
+        transcribe: int = self.special_tokens["[transcribe]"]
+
+        langs = tuple(LANGUAGES.keys())[: self.num_languages]
+        sot_sequence = [sot]
+        if self.language is not None:
+            sot_sequence.append(sot + 1 + langs.index(self.language))
+        if self.task is not None:
+            task_token: int = transcribe if self.task == "transcribe" else translate
+            sot_sequence.append(task_token)
+
+        self.sot_sequence = tuple(sot_sequence)
+
+    def get_vocab_size(self):
+        return self.encoding.n_vocab
+
+    def encode(self, text):
+        return self.encoding.encode(text, allowed_special="all")
+
+    def decode(self, token_ids: List[int], **kwargs) -> str:
+        return self.encoding.decode(token_ids, **kwargs)
+
+    @cached_property
+    def eot(self) -> int:
+        return self.encoding.eot_token
+
+    @cached_property
+    def stop(self) -> int:
+        return self.special_tokens["[STOP]"]
+
+    @cached_property
+    def start(self) -> int:
+        return self.special_tokens["[START]"]
+
+    @cached_property
+    def transcribe(self) -> int:
+        return self.special_tokens["[transcribe]"]
+
+    @cached_property
+    def translate(self) -> int:
+        return self.special_tokens["[translate]"]
+
+    @cached_property
+    def sot(self) -> int:
+        return self.special_tokens["[startoftranscript]"]
+
+    @cached_property
+    def sot_lm(self) -> int:
+        return self.special_tokens["[startoflm]"]
+
+    @cached_property
+    def sot_prev(self) -> int:
+        return self.special_tokens["[startofprev]"]
+
+    @cached_property
+    def no_speech(self) -> int:
+        return self.special_tokens["[nospeech]"]
+
+    @cached_property
+    def language_token(self) -> int:
+        """Returns the token id corresponding to the value of the `language` field"""
+        if self.language is None:
+            raise ValueError("This tokenizer does not have language token configured")
+
+        return self.to_language_token(self.language)
+
+    def to_language_token(self, language):
+        if token := self.special_tokens.get(f"[{language}]", None):
+            return token
+
+        raise KeyError(f"Language {language} not found in tokenizer.")
+
+    @cached_property
+    def all_language_tokens(self) -> Tuple[int]:
+        result = []
+        for token, token_id in self.special_tokens.items():
+            if token.strip("[]") in LANGUAGES:
+                result.append(token_id)
+        return tuple(result)[: self.num_languages]
+
+    @cached_property
+    def all_language_codes(self) -> Tuple[str]:
+        return tuple(self.decode([_l]).strip("[]") for _l in self.all_language_tokens)
+
+    @cached_property
+    def non_speech_tokens(self) -> Tuple[int]:
+        """
+        Returns the list of tokens to suppress in order to avoid any speaker tags or non-speech
+        annotations, to prevent sampling texts that are not actually spoken in the audio, e.g.
+
+        - ♪♪♪
+        - ( SPEAKING FOREIGN LANGUAGE )
+        - [DAVID] Hey there,
+
+        keeping basic punctuations like commas, periods, question marks, exclamation points, etc.
+        """
+        symbols = list('"#()*+/:;<=>@[\\]^_`{|}~「」『』')
+        symbols += (
+            "<< >> <<< >>> -- --- -( -[ (' (\" (( )) ((( ))) [[ ]] {{ }} ♪♪ ♪♪♪".split()
+        )
+
+        # symbols that may be a single token or multiple tokens depending on the tokenizer.
+        # In case they're multiple tokens, suppress the first token, which is safe because:
+        # These are between U+2640 and U+267F miscellaneous symbols that are okay to suppress
+        # in generations, and in the 3-byte UTF-8 representation they share the first two bytes.
+        miscellaneous = set("♩♪♫♬♭♮♯")
+        assert all(0x2640 <= ord(c) <= 0x267F for c in miscellaneous)
+
+        # allow hyphens "-" and single quotes "'" between words, but not at the beginning of a word
+        result = {self.encoding.encode(" -")[0], self.encoding.encode(" '")[0]}
+        for symbol in symbols + list(miscellaneous):
+            for tokens in [
+                self.encoding.encode(symbol),
+                self.encoding.encode(" " + symbol),
+            ]:
+                if len(tokens) == 1 or symbol in miscellaneous:
+                    result.add(tokens[0])
+
+        return tuple(sorted(result))
+
+    def split_to_word_tokens(self, tokens: List[int]):
+        if self.language in {"zh", "ja", "th", "lo", "my", "yue"}:
+            # These languages don't typically use spaces, so it is difficult to split words
+            # without morpheme analysis. Here, we instead split words at any
+            # position where the tokens are decoded as valid unicode points
+            return self.split_tokens_on_unicode(tokens)
+
+        return self.split_tokens_on_spaces(tokens)
+
+    def split_tokens_on_unicode(self, tokens: List[int]):
+        decoded_full = self.decode(tokens)
+        replacement_char = "\ufffd"
+
+        words = []
+        word_tokens = []
+        current_tokens = []
+        unicode_offset = 0
+
+        for token in tokens:
+            current_tokens.append(token)
+            decoded = self.decode(current_tokens)
+
+            if (
+                replacement_char not in decoded
+                or decoded_full[unicode_offset + decoded.index(replacement_char)]
+                == replacement_char
+            ):
+                words.append(decoded)
+                word_tokens.append(current_tokens)
+                current_tokens = []
+                unicode_offset += len(decoded)
+
+        return words, word_tokens
+
+    def split_tokens_on_spaces(self, tokens: List[int]):
+        subwords, subword_tokens_list = self.split_tokens_on_unicode(tokens)
+        words = []
+        word_tokens = []
+
+        for subword, subword_tokens in zip(subwords, subword_tokens_list):
+            special = subword_tokens[0] >= self.eot
+            with_space = subword.startswith(" ")
+            punctuation = subword.strip() in string.punctuation
+            if special or with_space or punctuation or len(words) == 0:
+                words.append(subword)
+                word_tokens.append(subword_tokens)
+            else:
+                words[-1] = words[-1] + subword
+                word_tokens[-1].extend(subword_tokens)
+
+        return words, word_tokens
+
+
+@lru_cache(maxsize=None)
+def get_encoding(name: str = "multilingual", num_languages: int = 100):
+    vocab_path = os.path.join(os.path.dirname(__file__), "assets", f"{name}.tiktoken")
+    ranks = {
+        base64.b64decode(token): int(rank)
+        for token, rank in (line.split() for line in open(vocab_path) if line)
+    }
+    n_vocab = len(ranks)
+    special_tokens = {}
+
+    specials = [
+        "[STOP]",
+        "[UNK]",
+        "[SPACE]",
+        "[START]",
+        "[nospk]",
+        "[spkemb]",
+        "[emotionemb]",
+        "[contextemb]",
+        "[sbreak]",
+        "[pbreak]",
+        "[uvbreak]",
+        "[bsing]",
+        "[esing]",
+        "[sing]",
+        "[hum]",
+        "[laugh]",
+        "[break]",
+        "[breath]",
+        "[oralsii]",
+        "[oralze]",
+        "[prolong]",
+        "[stress]",
+        "[bstrong]",
+        "[estrong]",
+        "[hiccup]",
+        "[inhale]",
+        "[exhale]",
+        "[emounknown]",
+        "[happy]",
+        "[neutral]",
+        "[sad]",
+        "[surprise]",
+        "[angry]",
+        "[disgust]",
+        "[emo]",
+        "[laugha]",
+        "[laughb]",
+        "[laughc]",
+        "[orala]",
+        "[oralb]",
+        "[oralc]",
+        "[orald]",
+        "[orale]",
+        "[breaka]",
+        "[breakb]",
+        "[breakc]",
+        "[breakd]",
+        "[breake]",
+        "[breakf]",
+        "[endoftext]",
+        "[startoftranscript]",
+        *[f"[{lang}]" for lang in list(LANGUAGES.keys())[:num_languages]],
+        "[translate]",
+        "[transcribe]",
+        "[startoflm]",
+        "[startofprev]",
+        "[nospeech]",
+    ]
+
+    for token in specials:
+        special_tokens[token] = n_vocab
+        n_vocab += 1
+
+    return tiktoken.Encoding(
+        name=os.path.basename(vocab_path),
+        explicit_n_vocab=n_vocab,
+        pat_str=r"""'s|'t|'re|'ve|'m|'ll|'d|\[[A-Z]+\]|\[[a-z]+\]|[\x{4e00}-\x{9df5}]| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""",
+        mergeable_ranks=ranks,
+        special_tokens=special_tokens,
+    )
+
+
+@lru_cache(maxsize=None)
+def get_tokenizer(
+    multilingual: bool,
+    *,
+    num_languages: int = 100,
+    language: Optional[str] = None,
+    task: Optional[str] = None,  # Literal["transcribe", "translate", None]
+) -> Tokenizer:
+    if language is not None:
+        language = language.lower()
+        if language not in LANGUAGES:
+            if language in TO_LANGUAGE_CODE:
+                language = TO_LANGUAGE_CODE[language]
+            else:
+                raise ValueError(f"Unsupported language: {language}")
+
+    if multilingual:
+        encoding_name = "multilingual"
+        language = language or "en"
+        task = task or "transcribe"
+    else:
+        encoding_name = "gpt2"
+        language = None
+        task = None
+
+    encoding = get_encoding(name=encoding_name, num_languages=num_languages)
+
+    return Tokenizer(
+        encoding=encoding, num_languages=num_languages, language=language, task=task
+    )

fireredtts/utils/utils.py

@@ -0,0 +1,37 @@
+import os
+from pathlib import Path
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+
+
+def load_audio(audiopath, sampling_rate):
+    """_summary_
+
+    Args:
+        audiopath (_type_): audio_path
+        sampling_rate (_type_): sampling_rate
+
+    Returns:
+        _type_: _description_
+    """
+    audio, lsr = torchaudio.load(audiopath)
+
+    # stereo to mono if needed
+    if audio.size(0) != 1:
+        audio = torch.mean(audio, dim=0, keepdim=True)
+
+    # resample
+    audio_resampled = torchaudio.functional.resample(audio, lsr, sampling_rate)
+    if torch.any(audio > 10) or not torch.any(audio < 0):
+        print(f"Error with {audiopath}. Max={audio.max()} min={audio.min()}")
+
+    if torch.any(audio_resampled > 10) or not torch.any(audio_resampled < 0):
+        print(
+            f"Error with {audiopath}. Max={audio_resampled.max()} min={audio_resampled.min()}"
+        )
+    # clip audio invalid values
+    audio.clip_(-1, 1)
+    audio_resampled.clip_(-1, 1)
+    return audio, lsr, audio_resampled

pretrained_models/README.md

@@ -0,0 +1,3 @@
+## Pretrained Models
+
+Download the required model files and place them in the folder `pretrained_models`