upload test model

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+text/custom_pypinyin_dict/__pycache__/cc_cedict_0.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text
+text/custom_pypinyin_dict/__pycache__/cc_cedict_0.cpython-38.pyc filter=lfs diff=lfs merge=lfs -text
+text/custom_pypinyin_dict/__pycache__/cc_cedict_1.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text
+text/custom_pypinyin_dict/__pycache__/cc_cedict_1.cpython-38.pyc filter=lfs diff=lfs merge=lfs -text
+text/custom_pypinyin_dict/__pycache__/cc_cedict_2.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text
+text/custom_pypinyin_dict/__pycache__/cc_cedict_2.cpython-38.pyc filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,138 @@
+import os
+
+from dataclasses import asdict
+from text import symbols
+import torch
+import torchaudio
+
+from utils.audio import LogMelSpectrogram
+from config import ModelConfig, VocosConfig, MelConfig
+from models.model import StableTTS
+from vocos_pytorch.models.model import Vocos
+from text.mandarin import chinese_to_cnm3
+from text import cleaned_text_to_sequence
+from datas.dataset import intersperse
+
+import gradio as gr
+import numpy as np
+import matplotlib.pyplot as plt
+from pathlib import Path
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+@ torch.inference_mode()
+def inference(text: str, ref_audio: torch.Tensor, checkpoint_path: str, step: int=10) -> torch.Tensor:
+    global last_checkpoint_path
+    if checkpoint_path != last_checkpoint_path:
+        tts_model.load_state_dict(torch.load(checkpoint_path, map_location='cpu')) 
+        last_checkpoint_path = checkpoint_path
+        
+    phonemizer = chinese_to_cnm3
+    
+    # prepare input for tts model
+    x = torch.tensor(intersperse(cleaned_text_to_sequence(phonemizer(text)), item=0), dtype=torch.long, device=device).unsqueeze(0)
+    x_len = torch.tensor([x.size(-1)], dtype=torch.long, device=device)
+    waveform, sr = torchaudio.load(ref_audio)
+    if sr != sample_rate:
+        waveform = torchaudio.functional.resample(waveform, sr, sample_rate)
+    y = mel_extractor(waveform).to(device)
+    
+    # inference
+    mel = tts_model.synthesise(x, x_len, step, y=y, temperature=1, length_scale=1)['decoder_outputs']
+    audio = vocoder(mel)
+    
+    # process output for gradio
+    audio_output = (sample_rate, (audio.cpu().squeeze(0).numpy() * 32767).astype(np.int16)) # (samplerate, int16 audio) for gr.Audio
+    mel_output = plot_mel_spectrogram(mel.cpu().squeeze(0).numpy()) # get the plot of mel
+    return audio_output, mel_output
+
+def get_pipeline(n_vocab: int, tts_model_config: ModelConfig, mel_config: MelConfig, vocoder_config: VocosConfig, tts_checkpoint_path, vocoder_checkpoint_path):
+    tts_model = StableTTS(n_vocab, mel_config.n_mels, **asdict(tts_model_config))
+    mel_extractor = LogMelSpectrogram(mel_config)
+    vocoder = Vocos(vocoder_config, mel_config)
+    # tts_model.load_state_dict(torch.load(tts_checkpoint_path, map_location='cpu'))
+    tts_model.to(device)
+    tts_model.eval()
+    vocoder.load_state_dict(torch.load(vocoder_checkpoint_path, map_location='cpu'))
+    vocoder.to(device)
+    vocoder.eval()
+    return tts_model, mel_extractor, vocoder
+
+def plot_mel_spectrogram(mel_spectrogram):
+    fig, ax = plt.subplots(figsize=(20, 8))
+    ax.imshow(mel_spectrogram, aspect='auto', origin='lower')
+    plt.axis('off')
+    fig.subplots_adjust(left=0, right=1, top=1, bottom=0) # remove white edges
+    return fig
+
+
+def main():
+    tts_model_config = ModelConfig()
+    mel_config = MelConfig()
+    vocoder_config = VocosConfig()
+
+    tts_checkpoint_path = './checkpoints' # the folder that contains StableTTS checkpoints
+    vocoder_checkpoint_path = './checkpoints/vocoder.pt'
+
+    global tts_model, mel_extractor, vocoder, sample_rate, last_checkpoint_path
+    sample_rate = mel_config.sample_rate
+    last_checkpoint_path = None
+    tts_model, mel_extractor, vocoder = get_pipeline(len(symbols), tts_model_config, mel_config, vocoder_config, tts_checkpoint_path, vocoder_checkpoint_path)
+    
+    tts_checkpoint_path = [path for path in Path(tts_checkpoint_path).rglob('*.pt') if 'optimizer' and 'vocoder' not in path.name]
+    audios = list(Path('./audios').rglob('*.wav'))
+
+    # gradio wabui
+    gui_title = 'StableTTS'
+    gui_description = """Next-generation TTS model using flow-matching and DiT, inspired by Stable Diffusion 3."""
+    with gr.Blocks(analytics_enabled=False) as demo:
+
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown(f"# {gui_title}")
+                gr.Markdown(gui_description)
+
+        with gr.Row():
+            with gr.Column():
+                input_text_gr = gr.Textbox(
+                    label="Input Text",
+                    info="One or two sentences at a time is better. Up to 200 text characters.",
+                    value="三国杀是一款风靡全球的以三国演义为背景的策略卡牌桌面游戏，经典新三国国战玩法，百万名将任由你搭配，楚雄争霸，等你决战沙场！",
+                )
+             
+                ref_audio_gr = gr.Dropdown(
+                    label='reference audio',
+                    choices=audios,
+                    value = 0
+                )
+                
+                
+                checkpoint_gr = gr.Dropdown(
+                    label='checkpoint',
+                    choices=tts_checkpoint_path,
+                    value = 0
+                )
+                
+                step_gr = gr.Slider(
+                    label='Step',
+                    minimum=1,
+                    maximum=100,
+                    value=25,
+                    step=1
+                )
+
+
+                tts_button = gr.Button("Send", elem_id="send-btn", visible=True)
+                
+            with gr.Column():
+                mel_gr = gr.Plot(label="Mel Visual")
+                audio_gr = gr.Audio(label="Synthesised Audio", autoplay=True)
+
+        tts_button.click(inference, [input_text_gr, ref_audio_gr, checkpoint_gr, step_gr], outputs=[audio_gr, mel_gr])
+
+    demo.queue()  
+    demo.launch(debug=True, show_api=True)
+
+
+if __name__ == '__main__':
+    main()

checkpoints/checkpoint_0.pt

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

checkpoints/vocoder.pt

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

config.py

@@ -0,0 +1,52 @@
+from dataclasses import dataclass
+
+@dataclass
+class MelConfig:
+    sample_rate: int = 44100
+    n_fft: int = 2048
+    win_length: int = 2048
+    hop_length: int = 512
+    f_min: float = 0.0
+    f_max: float = None
+    pad: int = 0
+    n_mels: int = 128
+    power: float = 1.0
+    normalized: bool = False
+    center: bool = False
+    pad_mode: str = "reflect"
+    mel_scale: str = "htk"
+    
+    def __post_init__(self):
+        if self.pad == 0:
+            self.pad = (self.n_fft - self.hop_length) // 2
+            
+@dataclass
+class ModelConfig:
+    hidden_channels: int = 192
+    filter_channels: int = 512
+    n_heads: int = 2
+    n_enc_layers: int = 3 
+    n_dec_layers: int = 2 
+    kernel_size: int = 3
+    p_dropout: int = 0.1
+    gin_channels: int = 192
+            
+@dataclass
+class TrainConfig:
+    train_dataset_path: str = 'filelists/filelist.json'
+    test_dataset_path: str = 'filelists/filelist.json'
+    batch_size: int = 52
+    learning_rate: float = 1e-4
+    num_epochs: int = 10000
+    model_save_path: str = './checkpoints'
+    log_dir: str = './runs'
+    log_interval: int = 128
+    save_interval: int = 15
+    warmup_steps: int = 200
+    
+@dataclass
+class VocosConfig:
+    input_channels: int = 128
+    dim: int = 512
+    intermediate_dim: int = 1536
+    num_layers: int = 8

datas/dataset.py

@@ -0,0 +1,52 @@
+import os
+
+import json
+import torch
+from torch.utils.data import Dataset
+
+from text import cleaned_text_to_sequence
+
+def intersperse(lst, item):
+  result = [item] * (len(lst) * 2 + 1)
+  result[1::2] = lst
+  return result
+    
+class StableDataset(Dataset):
+    def __init__(self, filelist_path, hop_length):
+        self.filelist_path = filelist_path     
+        self.hop_length = hop_length  
+        
+        self._load_filelist(filelist_path)
+
+    def _load_filelist(self, filelist_path):
+        filelist, lengths = [], []
+        with open(filelist_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                line = json.loads(line.strip())
+                filelist.append((line['mel_path'], line['phone']))
+                lengths.append(os.path.getsize(line['audio_path']) // (2 * self.hop_length))
+            
+        self.filelist = filelist
+        self.lengths = lengths
+    
+    def __len__(self):
+        return len(self.filelist)
+
+    def __getitem__(self, idx):
+        mel_path, phone = self.filelist[idx]
+        mel = torch.load(mel_path, map_location='cpu')
+        phone = torch.tensor(intersperse(cleaned_text_to_sequence(phone), 0), dtype=torch.long)
+        return mel, phone
+    
+def collate_fn(batch):
+    texts = [item[1] for item in batch]
+    mels = [item[0] for item in batch]
+    
+    text_lengths = torch.tensor([text.size(-1) for text in texts], dtype=torch.long)
+    mel_lengths = torch.tensor([mel.size(-1) for mel in mels], dtype=torch.long)
+    
+    # pad to the same length
+    texts_padded = torch.nested.to_padded_tensor(torch.nested.nested_tensor(texts), padding=0)
+    mels_padded = torch.nested.to_padded_tensor(torch.nested.nested_tensor(mels), padding=0)
+
+    return texts_padded, text_lengths, mels_padded, mel_lengths

datas/sampler.py

@@ -0,0 +1,121 @@
+import torch
+
+# reference: https://github.com/jaywalnut310/vits/blob/main/data_utils.py
+class DistributedBucketSampler(torch.utils.data.distributed.DistributedSampler):
+    """
+    Maintain similar input lengths in a batch.
+    Length groups are specified by boundaries.
+    Ex) boundaries = [b1, b2, b3] -> any batch is included either {x | b1 < length(x) <=b2} or {x | b2 < length(x) <= b3}.
+
+    It removes samples which are not included in the boundaries.
+    Ex) boundaries = [b1, b2, b3] -> any x s.t. length(x) <= b1 or length(x) > b3 are discarded.
+    """
+
+    def __init__(
+        self,
+        dataset,
+        batch_size,
+        boundaries,
+        num_replicas=None,
+        rank=None,
+        shuffle=True,
+    ):
+        super().__init__(dataset, num_replicas=num_replicas, rank=rank, shuffle=shuffle)
+        self.lengths = dataset.lengths
+        self.batch_size = batch_size
+        self.boundaries = boundaries
+
+        self.buckets, self.num_samples_per_bucket = self._create_buckets()
+        self.total_size = sum(self.num_samples_per_bucket)
+        self.num_samples = self.total_size // self.num_replicas
+
+    def _create_buckets(self):
+        buckets = [[] for _ in range(len(self.boundaries) - 1)]
+        for i in range(len(self.lengths)):
+            length = self.lengths[i]
+            idx_bucket = self._bisect(length)
+            if idx_bucket != -1:
+                buckets[idx_bucket].append(i)
+
+        for i in range(len(buckets) - 1, 0, -1):
+            # for i in range(len(buckets) - 1, -1, -1):
+            if len(buckets[i]) == 0:
+                buckets.pop(i)
+                self.boundaries.pop(i + 1)
+
+        num_samples_per_bucket = []
+        for i in range(len(buckets)):
+            len_bucket = len(buckets[i])
+            total_batch_size = self.num_replicas * self.batch_size
+            rem = (
+                total_batch_size - (len_bucket % total_batch_size)
+            ) % total_batch_size
+            num_samples_per_bucket.append(len_bucket + rem)
+        return buckets, num_samples_per_bucket
+
+    def __iter__(self):
+        # deterministically shuffle based on epoch
+        g = torch.Generator()
+        g.manual_seed(self.epoch)
+
+        indices = []
+        if self.shuffle:
+            for bucket in self.buckets:
+                indices.append(torch.randperm(len(bucket), generator=g).tolist())
+        else:
+            for bucket in self.buckets:
+                indices.append(list(range(len(bucket))))
+
+        batches = []
+        for i in range(len(self.buckets)):
+            bucket = self.buckets[i]
+            len_bucket = len(bucket)
+            ids_bucket = indices[i]
+            num_samples_bucket = self.num_samples_per_bucket[i]
+
+            # add extra samples to make it evenly divisible
+            rem = num_samples_bucket - len_bucket
+            ids_bucket = (
+                ids_bucket
+                + ids_bucket * (rem // len_bucket)
+                + ids_bucket[: (rem % len_bucket)]
+            )
+
+            # subsample
+            ids_bucket = ids_bucket[self.rank :: self.num_replicas]
+
+            # batching
+            for j in range(len(ids_bucket) // self.batch_size):
+                batch = [
+                    bucket[idx]
+                    for idx in ids_bucket[
+                        j * self.batch_size : (j + 1) * self.batch_size
+                    ]
+                ]
+                batches.append(batch)
+
+        if self.shuffle:
+            batch_ids = torch.randperm(len(batches), generator=g).tolist()
+            batches = [batches[i] for i in batch_ids]
+        self.batches = batches
+
+        assert len(self.batches) * self.batch_size == self.num_samples
+        return iter(self.batches)
+
+    def _bisect(self, x, lo=0, hi=None):
+        if hi is None:
+            hi = len(self.boundaries) - 1
+
+        if hi > lo:
+            mid = (hi + lo) // 2
+            if self.boundaries[mid] < x and x <= self.boundaries[mid + 1]:
+                return mid
+            elif x <= self.boundaries[mid]:
+                return self._bisect(x, lo, mid)
+            else:
+                return self._bisect(x, mid + 1, hi)
+        else:
+            return -1
+
+    def __len__(self):
+        return self.num_samples // self.batch_size

models/dit.py

@@ -0,0 +1,205 @@
+# References:
+# https://github.com/shivammehta25/Matcha-TTS/blob/main/matcha/models/components/transformer.py
+# https://github.com/jaywalnut310/vits/blob/main/attentions.py
+# https://github.com/pytorch-labs/gpt-fast/blob/main/model.py
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class FFN(nn.Module):
+  def __init__(self, in_channels, out_channels, filter_channels, kernel_size, p_dropout=0., gin_channels=0):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.filter_channels = filter_channels
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+    self.gin_channels = gin_channels
+    
+    self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size // 2)
+    self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size, padding=kernel_size // 2)
+    self.drop = nn.Dropout(p_dropout)
+    self.act1 = nn.GELU(approximate="tanh")
+
+  def forward(self, x, x_mask):
+        x = self.conv_1(x * x_mask)
+        x = self.act1(x)
+        x = self.drop(x)
+        x = self.conv_2(x * x_mask)
+        return x * x_mask
+    
+class MultiHeadAttention(nn.Module):
+  def __init__(self, channels, out_channels, n_heads, p_dropout=0.):
+    super().__init__()
+    assert channels % n_heads == 0
+
+    self.channels = channels
+    self.out_channels = out_channels
+    self.n_heads = n_heads
+    self.p_dropout = p_dropout
+
+    self.k_channels = channels // n_heads
+    self.conv_q = torch.nn.Conv1d(channels, channels, 1)
+    self.conv_k = torch.nn.Conv1d(channels, channels, 1)
+    self.conv_v = torch.nn.Conv1d(channels, channels, 1)
+
+    # from https://nn.labml.ai/transformers/rope/index.html
+    self.query_rotary_pe = RotaryPositionalEmbeddings(self.k_channels * 0.5)
+    self.key_rotary_pe = RotaryPositionalEmbeddings(self.k_channels * 0.5)
+
+    self.conv_o = torch.nn.Conv1d(channels, out_channels, 1)
+    self.drop = torch.nn.Dropout(p_dropout)
+
+    torch.nn.init.xavier_uniform_(self.conv_q.weight)
+    torch.nn.init.xavier_uniform_(self.conv_k.weight)
+    torch.nn.init.xavier_uniform_(self.conv_v.weight)
+
+  def forward(self, x, attn_mask=None):
+      q = self.conv_q(x)
+      k = self.conv_k(x)
+      v = self.conv_v(x)
+
+      x = self.attention(q, k, v, mask=attn_mask)
+
+      x = self.conv_o(x)
+      return x
+
+  def attention(self, query, key, value, mask=None):
+      b, d, t_s, t_t = (*key.size(), query.size(2))
+      query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
+      key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+      value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+
+      query = self.query_rotary_pe(query) # [b, n_head, t, c // n_head]
+      key = self.key_rotary_pe(key)
+
+      output = F.scaled_dot_product_attention(query, key, value, attn_mask=mask, dropout_p=self.p_dropout if self.training else 0)
+      output = output.transpose(2, 3).contiguous().view(b, d, t_t)  # [b, n_h, t_t, d_k] -> [b, d, t_t]
+      return output
+
+# modified from https://github.com/sh-lee-prml/HierSpeechpp/blob/main/modules.py#L390    
+class DiTConVBlock(nn.Module):
+    """
+    A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+    """
+    def __init__(self, hidden_channels, filter_channels, num_heads, kernel_size=3, p_dropout=0.1, gin_channels=0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_channels, elementwise_affine=False, eps=1e-6)
+        self.attn = MultiHeadAttention(hidden_channels, hidden_channels, num_heads, p_dropout)
+        self.norm2 = nn.LayerNorm(hidden_channels, elementwise_affine=False, eps=1e-6)
+        self.mlp = FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout)
+        self.adaLN_modulation = nn.Sequential(
+            nn.Linear(gin_channels, hidden_channels) if gin_channels != hidden_channels else nn.Identity(),
+            nn.SiLU(),
+            nn.Linear(hidden_channels, 6 * hidden_channels, bias=True)
+        )
+            
+    def forward(self, x, c, x_mask):
+        """
+        Args:
+            x : [batch_size, channel, time]
+            c : [batch_size, channel]
+            x_mask : [batch_size, 1, time]
+        return the same shape as x
+        """
+        x = x * x_mask
+        attn_mask = x_mask.unsqueeze(1) * x_mask.unsqueeze(-1) # shape: [batch_size, 1, time, time]
+        # attn_mask = attn_mask.to(torch.bool)
+        
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).unsqueeze(2).chunk(6, dim=1) # shape: [batch_size, channel, 1]
+        x = x + gate_msa * self.attn(self.modulate(self.norm1(x.transpose(1,2)).transpose(1,2), shift_msa, scale_msa), attn_mask) * x_mask
+        x = x + gate_mlp * self.mlp(self.modulate(self.norm2(x.transpose(1,2)).transpose(1,2), shift_mlp, scale_mlp), x_mask)
+        
+        # no condition version
+        # x = x + self.attn(self.norm1(x.transpose(1,2)).transpose(1,2),  attn_mask)
+        # x = x + self.mlp(self.norm1(x.transpose(1,2)).transpose(1,2), x_mask)
+        return x
+    
+    @staticmethod
+    def modulate(x, shift, scale):
+        return x * (1 + scale) + shift
+  
+class RotaryPositionalEmbeddings(nn.Module):
+    """
+    ## RoPE module
+
+    Rotary encoding transforms pairs of features by rotating in the 2D plane.
+    That is, it organizes the $d$ features as $\frac{d}{2}$ pairs.
+    Each pair can be considered a coordinate in a 2D plane, and the encoding will rotate it
+    by an angle depending on the position of the token.
+    """
+
+    def __init__(self, d: int, base: int = 10_000):
+        r"""
+        * `d` is the number of features $d$
+        * `base` is the constant used for calculating $\Theta$
+        """
+        super().__init__()
+
+        self.base = base
+        self.d = int(d)
+        self.cos_cached = None
+        self.sin_cached = None
+
+    def _build_cache(self, x: torch.Tensor):
+        r"""
+        Cache $\cos$ and $\sin$ values
+        """
+        # Return if cache is already built
+        if self.cos_cached is not None and x.shape[0] <= self.cos_cached.shape[0]:
+            return
+
+        # Get sequence length
+        seq_len = x.shape[0]
+
+        # $\Theta = {\theta_i = 10000^{-\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
+        theta = 1.0 / (self.base ** (torch.arange(0, self.d, 2).float() / self.d)).to(x.device)
+
+        # Create position indexes `[0, 1, ..., seq_len - 1]`
+        seq_idx = torch.arange(seq_len, device=x.device).float().to(x.device)
+
+        # Calculate the product of position index and $\theta_i$
+        idx_theta = torch.einsum("n,d->nd", seq_idx, theta)
+
+        # Concatenate so that for row $m$ we have
+        # $[m \theta_0, m \theta_1, ..., m \theta_{\frac{d}{2}}, m \theta_0, m \theta_1, ..., m \theta_{\frac{d}{2}}]$
+        idx_theta2 = torch.cat([idx_theta, idx_theta], dim=1)
+
+        # Cache them
+        self.cos_cached = idx_theta2.cos()[:, None, None, :]
+        self.sin_cached = idx_theta2.sin()[:, None, None, :]
+
+    def _neg_half(self, x: torch.Tensor):
+        # $\frac{d}{2}$
+        d_2 = self.d // 2
+
+        # Calculate $[-x^{(\frac{d}{2} + 1)}, -x^{(\frac{d}{2} + 2)}, ..., -x^{(d)}, x^{(1)}, x^{(2)}, ..., x^{(\frac{d}{2})}]$
+        return torch.cat([-x[:, :, :, d_2:], x[:, :, :, :d_2]], dim=-1)
+
+    def forward(self, x: torch.Tensor):
+        """
+        * `x` is the Tensor at the head of a key or a query with shape `[seq_len, batch_size, n_heads, d]`
+        """
+        # Cache $\cos$ and $\sin$ values
+        x = x.permute(2, 0, 1, 3) # b h t d -> t b h d
+
+        self._build_cache(x)
+
+        # Split the features, we can choose to apply rotary embeddings only to a partial set of features.
+        x_rope, x_pass = x[..., : self.d], x[..., self.d :]
+
+        # Calculate
+        # $[-x^{(\frac{d}{2} + 1)}, -x^{(\frac{d}{2} + 2)}, ..., -x^{(d)}, x^{(1)}, x^{(2)}, ..., x^{(\frac{d}{2})}]$
+        neg_half_x = self._neg_half(x_rope)
+
+        x_rope = (x_rope * self.cos_cached[: x.shape[0]]) + (neg_half_x * self.sin_cached[: x.shape[0]])
+
+        return torch.cat((x_rope, x_pass), dim=-1).permute(1, 2, 0, 3) # t b h d -> b h t d
+
+class Transpose(nn.Identity):
+    """(N, T, D) -> (N, D, T)"""
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return input.transpose(1, 2)
+    

models/duration_predictor.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn as nn
+
+# modified from https://github.com/jaywalnut310/vits/blob/main/models.py#L98
+class DurationPredictor(nn.Module):
+  def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0):
+    super().__init__()
+
+    self.in_channels = in_channels
+    self.filter_channels = filter_channels
+    self.kernel_size = kernel_size
+    self.p_dropout = p_dropout
+    self.gin_channels = gin_channels
+
+    self.drop = nn.Dropout(p_dropout)
+    self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size//2)
+    self.norm_1 = nn.LayerNorm(filter_channels)
+    self.conv_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size//2)
+    self.norm_2 = nn.LayerNorm(filter_channels)
+    self.proj = nn.Conv1d(filter_channels, 1, 1)
+
+    self.cond = nn.Conv1d(gin_channels, in_channels, 1)
+
+  def forward(self, x, x_mask, g):
+    x = x.detach()
+    x = x + self.cond(g.unsqueeze(2).detach())
+    x = self.conv_1(x * x_mask)
+    x = torch.relu(x)
+    x = self.norm_1(x.transpose(1,2)).transpose(1,2)
+    x = self.drop(x)
+    x = self.conv_2(x * x_mask)
+    x = torch.relu(x)
+    x = self.norm_2(x.transpose(1,2)).transpose(1,2)
+    x = self.drop(x)
+    x = self.proj(x * x_mask)
+    return x * x_mask
+  
+def duration_loss(logw, logw_, lengths):
+    loss = torch.sum((logw - logw_) ** 2) / torch.sum(lengths)
+    return loss

models/estimator.py

@@ -0,0 +1,161 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from models.dit import DiTConVBlock
+
+class DitWrapper(nn.Module):
+    """ add FiLM layer to condition time embedding to DiT """
+    def __init__(self, hidden_channels, filter_channels, num_heads, kernel_size=3, p_dropout=0.1, gin_channels=0, time_channels=0):
+        super().__init__()
+        self.time_fusion = FiLMLayer(hidden_channels, time_channels)
+        self.conv1 = ConvNeXtBlock(hidden_channels, filter_channels, gin_channels)
+        self.conv2 = ConvNeXtBlock(hidden_channels, filter_channels, gin_channels)
+        self.conv3 = ConvNeXtBlock(hidden_channels, filter_channels, gin_channels)
+        self.block = DiTConVBlock(hidden_channels, hidden_channels, num_heads, kernel_size, p_dropout, gin_channels)
+            
+    def forward(self, x, c, t, x_mask):
+        x = self.time_fusion(x, t) * x_mask
+        x = self.conv1(x, c, x_mask)
+        x = self.conv2(x, c, x_mask)
+        x = self.conv3(x, c, x_mask)
+        x = self.block(x, c, x_mask)
+        return x
+
+class FiLMLayer(nn.Module):
+    """
+    Feature-wise Linear Modulation (FiLM) layer
+    Reference: https://arxiv.org/abs/1709.07871
+    """
+    def __init__(self, in_channels, cond_channels):
+
+        super(FiLMLayer, self).__init__()
+        self.in_channels = in_channels
+        self.film = nn.Conv1d(cond_channels, in_channels * 2, 1)
+
+    def forward(self, x, c):
+        gamma, beta = torch.chunk(self.film(c.unsqueeze(2)), chunks=2, dim=1)
+        return gamma * x + beta
+    
+class ConvNeXtBlock(nn.Module):
+    def __init__(self, in_channels, filter_channels, gin_channels):
+        super().__init__()
+        self.dwconv = nn.Conv1d(in_channels, in_channels, kernel_size=7, padding=3, groups=in_channels)
+        self.norm = StyleAdaptiveLayerNorm(in_channels, gin_channels)
+        self.pwconv = nn.Sequential(nn.Linear(in_channels, filter_channels),
+                                    nn.GELU(),
+                                    nn.Linear(filter_channels, in_channels))
+
+    def forward(self, x, c, x_mask) -> torch.Tensor:
+        residual = x
+        x = self.dwconv(x) * x_mask
+        x = self.norm(x.transpose(1, 2), c)
+        x = self.pwconv(x).transpose(1, 2)
+        x = residual + x
+        return x * x_mask
+
+class StyleAdaptiveLayerNorm(nn.Module):
+    def __init__(self, in_channels, cond_channels):
+        """
+        Style Adaptive Layer Normalization (SALN) module.
+
+        Parameters:
+        in_channels: The number of channels in the input feature maps.
+        cond_channels: The number of channels in the conditioning input.
+        """
+        super(StyleAdaptiveLayerNorm, self).__init__()
+        self.in_channels = in_channels
+
+        self.saln = nn.Linear(cond_channels, in_channels * 2, 1)
+        self.norm = nn.LayerNorm(in_channels, elementwise_affine=False)
+        
+        self.reset_parameters()
+        
+    def reset_parameters(self):
+        nn.init.constant_(self.saln.bias.data[:self.in_channels], 1)
+        nn.init.constant_(self.saln.bias.data[self.in_channels:], 0)
+
+    def forward(self, x, c):
+        gamma, beta = torch.chunk(self.saln(c.unsqueeze(1)), chunks=2, dim=-1)
+        return gamma * self.norm(x) + beta
+        
+    
+class SinusoidalPosEmb(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)
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=x.device).float() * -emb)
+        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+class TimestepEmbedding(nn.Module):
+    def __init__(self, in_channels, out_channels, filter_channels):
+        super().__init__()
+
+        self.layer = nn.Sequential(
+            nn.Linear(in_channels, filter_channels),
+            nn.SiLU(inplace=True),
+            nn.Linear(filter_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.layer(x)
+
+# reference: https://github.com/shivammehta25/Matcha-TTS/blob/main/matcha/models/components/decoder.py
+class Decoder(nn.Module):
+    def __init__(self, hidden_channels, out_channels, filter_channels, dropout=0.05, n_layers=1, n_heads=4, kernel_size=3, gin_channels=0):
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.filter_channels = filter_channels
+
+        self.time_embeddings = SinusoidalPosEmb(hidden_channels)
+        self.time_mlp = TimestepEmbedding(hidden_channels, hidden_channels, filter_channels)
+
+        
+        self.blocks = nn.ModuleList([DitWrapper(hidden_channels, filter_channels, n_heads, kernel_size, dropout, gin_channels, hidden_channels) for _ in range(n_layers)])
+        self.final_proj = nn.Conv1d(hidden_channels, out_channels, 1)
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        for block in self.blocks:
+            nn.init.constant_(block.block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.block.adaLN_modulation[-1].bias, 0)
+
+    def forward(self, x, mask, mu, t, c):
+        """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)
+            c (_type_): shape (batch_size, gin_channels)
+
+        Raises:
+            ValueError: _description_
+            ValueError: _description_
+
+        Returns:
+            _type_: _description_
+        """
+
+        t = self.time_mlp(self.time_embeddings(t))
+        x = torch.cat((x, mu), dim=1)
+
+        for block in self.blocks:
+            x = block(x, c, t, mask)
+
+        output = self.final_proj(x * mask)
+
+        return output * mask

models/flow_matching.py

@@ -0,0 +1,108 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from models.estimator import Decoder
+
+# copied from https://github.com/jaywalnut310/vits/blob/main/commons.py#L121
+def sequence_mask(length: torch.Tensor, max_length: int = None) -> torch.Tensor:
+    if max_length is None:
+        max_length = length.max()
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+# modified from https://github.com/shivammehta25/Matcha-TTS/blob/main/matcha/models/components/flow_matching.py
+class CFMDecoder(torch.nn.Module):
+    def __init__(self, hidden_channels, out_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, gin_channels):
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.filter_channels = filter_channels
+        self.gin_channels = gin_channels
+        self.sigma_min = 1e-4
+
+        self.estimator = Decoder(hidden_channels, out_channels, filter_channels, p_dropout, n_layers, n_heads, kernel_size, gin_channels)
+
+    @torch.inference_mode()
+    def forward(self, mu, mask, n_timesteps, temperature=1.0, c=None):
+        """Forward diffusion
+
+        Args:
+            mu (torch.Tensor): output of encoder
+                shape: (batch_size, n_feats, mel_timesteps)
+            mask (torch.Tensor): output_mask
+                shape: (batch_size, 1, mel_timesteps)
+            n_timesteps (int): number of diffusion steps
+            temperature (float, optional): temperature for scaling noise. Defaults to 1.0.
+            c (torch.Tensor, optional): shape: (batch_size, gin_channels)
+
+        Returns:
+            sample: generated mel-spectrogram
+                shape: (batch_size, n_feats, mel_timesteps)
+        """
+        z = torch.randn_like(mu) * temperature
+        t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device)
+        return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, c=c)
+
+    def solve_euler(self, x, t_span, mu, mask, c):
+        """
+        Fixed euler solver for ODEs.
+        Args:
+            x (torch.Tensor): random noise
+            t_span (torch.Tensor): n_timesteps interpolated
+                shape: (n_timesteps + 1,)
+            mu (torch.Tensor): output of encoder
+                shape: (batch_size, n_feats, mel_timesteps)
+            mask (torch.Tensor): output_mask
+                shape: (batch_size, 1, mel_timesteps)
+            c (torch.Tensor, optional): speaker condition.
+                shape: (batch_size, gin_channels)
+        """
+        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, c)
+
+            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 compute_loss(self, x1, mask, mu, c):
+        """Computes diffusion loss
+
+        Args:
+            x1 (torch.Tensor): Target
+                shape: (batch_size, n_feats, mel_timesteps)
+            mask (torch.Tensor): target mask
+                shape: (batch_size, 1, mel_timesteps)
+            mu (torch.Tensor): output of encoder
+                shape: (batch_size, n_feats, mel_timesteps)
+            c (torch.Tensor, optional): speaker condition.
+
+        Returns:
+            loss: conditional flow matching loss
+            y: conditional flow
+                shape: (batch_size, n_feats, mel_timesteps)
+        """
+        b, _, t = mu.shape
+
+        # random timestep
+        t = torch.rand([b, 1, 1], device=mu.device, dtype=mu.dtype)
+        # sample noise p(x_0)
+        z = torch.randn_like(x1)
+
+        y = (1 - (1 - self.sigma_min) * t) * z + t * x1
+        u = x1 - (1 - self.sigma_min) * z
+
+        loss = F.mse_loss(self.estimator(y, mask, mu, t.squeeze(), c), u, reduction="sum") / (
+            torch.sum(mask) * u.shape[1]
+        )
+        return loss, y

models/model.py

@@ -0,0 +1,194 @@
+import math
+import torch
+import torch.nn as nn
+
+import monotonic_align
+from models.text_encoder import TextEncoder
+from models.flow_matching import CFMDecoder
+from models.reference_encoder import MelStyleEncoder
+from models.duration_predictor import DurationPredictor, duration_loss
+
+def sequence_mask(length: torch.Tensor, max_length: int = None) -> torch.Tensor:
+    if max_length is None:
+        max_length = length.max()
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+def convert_pad_shape(pad_shape):
+    inverted_shape = pad_shape[::-1]
+    pad_shape = [item for sublist in inverted_shape for item in sublist]
+    return pad_shape
+
+def generate_path(duration, mask):
+    device = duration.device
+
+    b, t_x, t_y = mask.shape
+    cum_duration = torch.cumsum(duration, 1)
+    path = torch.zeros(b, t_x, t_y, dtype=mask.dtype).to(device=device)
+
+    cum_duration_flat = cum_duration.view(b * t_x)
+    path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype)
+    path = path.view(b, t_x, t_y)
+    path = path - torch.nn.functional.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
+    path = path * mask
+    return path
+
+# modified from https://github.com/shivammehta25/Matcha-TTS/blob/main/matcha/models/matcha_tts.py
+class StableTTS(nn.Module):
+    def __init__(self, n_vocab, mel_channels, hidden_channels, filter_channels, n_heads, n_enc_layers, n_dec_layers, kernel_size, p_dropout, gin_channels):
+        super().__init__()
+
+        self.n_vocab = n_vocab
+        self.mel_channels = mel_channels
+
+        self.encoder = TextEncoder(n_vocab, mel_channels, hidden_channels, filter_channels, n_heads, n_enc_layers, kernel_size, p_dropout, gin_channels)
+        self.ref_encoder = MelStyleEncoder(mel_channels, style_vector_dim=gin_channels, style_kernel_size=3)
+        self.dp = DurationPredictor(hidden_channels, filter_channels, kernel_size, p_dropout, gin_channels)
+        self.decoder = CFMDecoder(mel_channels + mel_channels, mel_channels, filter_channels, n_heads, n_dec_layers, kernel_size, p_dropout, gin_channels)
+
+    @torch.inference_mode()
+    def synthesise(self, x, x_lengths, n_timesteps, temperature=1.0, y=None, length_scale=1.0):
+        """
+        Generates mel-spectrogram from text. Returns:
+            1. encoder outputs
+            2. decoder outputs
+            3. generated alignment
+
+        Args:
+            x (torch.Tensor): batch of texts, converted to a tensor with phoneme embedding ids.
+                shape: (batch_size, max_text_length)
+            x_lengths (torch.Tensor): lengths of texts in batch.
+                shape: (batch_size,)
+            n_timesteps (int): number of steps to use for reverse diffusion in decoder.
+            temperature (float, optional): controls variance of terminal distribution.
+            y (torch.Tensor): mel spectrogram of reference audio
+                shape: (batch_size, mel_channels, time)
+            length_scale (float, optional): controls speech pace.
+                Increase value to slow down generated speech and vice versa.
+
+        Returns:
+            dict: {
+                "encoder_outputs": torch.Tensor, shape: (batch_size, n_feats, max_mel_length),
+                # Average mel spectrogram generated by the encoder
+                "decoder_outputs": torch.Tensor, shape: (batch_size, n_feats, max_mel_length),
+                # Refined mel spectrogram improved by the CFM
+                "attn": torch.Tensor, shape: (batch_size, max_text_length, max_mel_length),
+                # Alignment map between text and mel spectrogram
+        """
+
+        # Get encoder_outputs `mu_x` and log-scaled token durations `logw`
+        c = self.ref_encoder(y, None)
+        x, mu_x, x_mask = self.encoder(x, c, x_lengths)
+        logw = self.dp(x, x_mask, c)
+
+        w = torch.exp(logw) * x_mask
+        w_ceil = torch.ceil(w) * length_scale
+        y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
+        y_max_length = y_lengths.max()
+
+        # Using obtained durations `w` construct alignment map `attn`
+        y_mask = sequence_mask(y_lengths, y_max_length).unsqueeze(1).to(x_mask.dtype)
+        attn_mask = x_mask.unsqueeze(-1) * y_mask.unsqueeze(2)
+        attn = generate_path(w_ceil.squeeze(1), attn_mask.squeeze(1)).unsqueeze(1)
+
+        # Align encoded text and get mu_y
+        mu_y = torch.matmul(attn.squeeze(1).transpose(1, 2), mu_x.transpose(1, 2))
+        mu_y = mu_y.transpose(1, 2)
+        encoder_outputs = mu_y[:, :, :y_max_length]
+
+        # Generate sample tracing the probability flow
+        decoder_outputs = self.decoder(mu_y, y_mask, n_timesteps, temperature, c)
+        decoder_outputs = decoder_outputs[:, :, :y_max_length]
+
+
+        return {
+            "encoder_outputs": encoder_outputs,
+            "decoder_outputs": decoder_outputs,
+            "attn": attn[:, :, :y_max_length],
+        }
+
+    def forward(self, x, x_lengths, y, y_lengths):
+        """
+        Computes 3 losses:
+            1. duration loss: loss between predicted token durations and those extracted by Monotinic Alignment Search (MAS).
+            2. prior loss: loss between mel-spectrogram and encoder outputs.
+            3. flow matching loss: loss between mel-spectrogram and decoder outputs.
+
+        Args:
+            x (torch.Tensor): batch of texts, converted to a tensor with phoneme embedding ids.
+                shape: (batch_size, max_text_length)
+            x_lengths (torch.Tensor): lengths of texts in batch.
+                shape: (batch_size,)
+            y (torch.Tensor): batch of corresponding mel-spectrograms.
+                shape: (batch_size, n_feats, max_mel_length)
+            y_lengths (torch.Tensor): lengths of mel-spectrograms in batch.
+                shape: (batch_size,)
+        """
+        # Get encoder_outputs `mu_x` and log-scaled token durations `logw`
+        y_mask = sequence_mask(y_lengths, y.size(2)).unsqueeze(1).to(y.dtype)
+        c = self.ref_encoder(y, y_mask)
+        
+        x, mu_x, x_mask = self.encoder(x, c, x_lengths)
+        logw = self.dp(x, x_mask, c)
+        
+        attn_mask = x_mask.unsqueeze(-1) * y_mask.unsqueeze(2)
+        
+        # Use MAS to find most likely alignment `attn` between text and mel-spectrogram
+        
+        # I'm not sure why the MAS code in Matcha TTS and Grad TTS could not align in StableTTS
+        # so I use the code from https://github.com/p0p4k/pflowtts_pytorch/blob/master/pflow/models/pflow_tts.py and it works
+        # Welcome everyone to solve this problem QAQ
+        
+        with torch.no_grad():
+            # const = -0.5 * math.log(2 * math.pi) * self.n_feats
+            # const = -0.5 * math.log(2 * math.pi) * self.mel_channels
+            # factor = -0.5 * torch.ones(mu_x.shape, dtype=mu_x.dtype, device=mu_x.device)
+            # y_square = torch.matmul(factor.transpose(1, 2), y**2)
+            # y_mu_double = torch.matmul(2.0 * (factor * mu_x).transpose(1, 2), y)
+            # mu_square = torch.sum(factor * (mu_x**2), 1).unsqueeze(-1)
+            # log_prior = y_square - y_mu_double + mu_square + const
+            
+            s_p_sq_r = torch.ones_like(mu_x) # [b, d, t]
+            # s_p_sq_r = torch.exp(-2 * logx) 
+            neg_cent1 = torch.sum(
+                -0.5 * math.log(2 * math.pi)- torch.zeros_like(mu_x), [1], keepdim=True
+            )
+            # neg_cent1 = torch.sum(
+            #     -0.5 * math.log(2 * math.pi) - logx, [1], keepdim=True
+            #     ) # [b, 1, t_s]
+            neg_cent2 = torch.einsum("bdt, bds -> bts", -0.5 * (y**2), s_p_sq_r)
+            neg_cent3 = torch.einsum("bdt, bds -> bts", y, (mu_x * s_p_sq_r))
+            neg_cent4 = torch.sum(
+                -0.5 * (mu_x**2) * s_p_sq_r, [1], keepdim=True
+            )  
+            neg_cent = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4
+            
+            attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
+            
+            attn = (
+                 monotonic_align.maximum_path(neg_cent, attn_mask.squeeze(1)).unsqueeze(1).detach()
+            )
+
+            # attn = monotonic_align.maximum_path(neg_cent, attn_mask.squeeze(1))
+            # attn = attn.detach()
+
+        # Compute loss between predicted log-scaled durations and those obtained from MAS
+        # refered to as prior loss in the paper
+        logw_ = torch.log(1e-8 + attn.sum(2)) * x_mask
+        # logw_ = torch.log(1e-8 + torch.sum(attn.unsqueeze(1), -1)) * x_mask
+        dur_loss = duration_loss(logw, logw_, x_lengths)
+
+
+        # Align encoded text with mel-spectrogram and get mu_y segment
+        attn = attn.squeeze(1).transpose(1,2)
+        mu_y = torch.matmul(attn.squeeze(1).transpose(1, 2), mu_x.transpose(1, 2))
+        mu_y = mu_y.transpose(1, 2)
+
+        # Compute loss of the decoder
+        diff_loss, _ = self.decoder.compute_loss(y, y_mask, mu_y, c)
+        # diff_loss = torch.tensor([0], device=mu_y.device)
+        
+        prior_loss = torch.sum(0.5 * ((y - mu_y) ** 2 + math.log(2 * math.pi)) * y_mask)
+        prior_loss = prior_loss / (torch.sum(y_mask) * self.mel_channels)
+
+        return dur_loss, diff_loss, prior_loss, attn

models/reference_encoder.py

@@ -0,0 +1,93 @@
+import torch
+import torch.nn as nn
+    
+class Conv1dGLU(nn.Module):
+    """
+    Conv1d + GLU(Gated Linear Unit) with residual connection.
+    For GLU refer to https://arxiv.org/abs/1612.08083 paper.
+    """
+
+    def __init__(self, in_channels, out_channels, kernel_size, dropout):
+        super(Conv1dGLU, self).__init__()
+        self.out_channels = out_channels
+        self.conv1 = nn.Conv1d(in_channels, 2 * out_channels, kernel_size=kernel_size, padding=kernel_size // 2)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        residual = x
+        x = self.conv1(x)
+        x1, x2 = torch.split(x, self.out_channels, dim=1)
+        x = x1 * torch.sigmoid(x2)
+        x = residual + self.dropout(x)
+        return x
+
+# modified from https://github.com/RVC-Boss/GPT-SoVITS/blob/main/GPT_SoVITS/module/modules.py#L766    
+class MelStyleEncoder(nn.Module):
+    """MelStyleEncoder"""
+
+    def __init__(
+        self,
+        n_mel_channels=80,
+        style_hidden=128,
+        style_vector_dim=256,
+        style_kernel_size=5,
+        style_head=2,
+        dropout=0.1,
+    ):
+        super(MelStyleEncoder, self).__init__()
+        self.in_dim = n_mel_channels
+        self.hidden_dim = style_hidden
+        self.out_dim = style_vector_dim
+        self.kernel_size = style_kernel_size
+        self.n_head = style_head
+        self.dropout = dropout
+
+        self.spectral = nn.Sequential(
+            nn.Linear(self.in_dim, self.hidden_dim),
+            nn.Mish(inplace=True),
+            nn.Dropout(self.dropout),
+            nn.Linear(self.hidden_dim, self.hidden_dim),
+            nn.Mish(inplace=True),
+            nn.Dropout(self.dropout),
+        )
+
+        self.temporal = nn.Sequential(
+            Conv1dGLU(self.hidden_dim, self.hidden_dim, self.kernel_size, self.dropout),
+            Conv1dGLU(self.hidden_dim, self.hidden_dim, self.kernel_size, self.dropout),
+        )
+
+        self.slf_attn = nn.MultiheadAttention(
+            self.hidden_dim,
+            self.n_head,
+            self.dropout,
+            batch_first=True
+        )
+
+        self.fc = nn.Linear(self.hidden_dim, self.out_dim)
+
+    def temporal_avg_pool(self, x, mask=None):
+        if mask is None:
+            return torch.mean(x, dim=1)
+        else:
+            len_ = (~mask).sum(dim=1).unsqueeze(1).type_as(x)
+            return torch.sum(x * ~mask.unsqueeze(-1), dim=1) / len_
+
+    def forward(self, x, x_mask=None):
+        x = x.transpose(1, 2)
+
+        # spectral
+        x = self.spectral(x)
+        # temporal
+        x = x.transpose(1, 2)
+        x = self.temporal(x)
+        x = x.transpose(1, 2)
+        # self-attention
+        if x_mask is not None:
+            x_mask = ~x_mask.squeeze(1).to(torch.bool)   
+        x, _ = self.slf_attn(x, x, x, key_padding_mask=x_mask)
+        # fc
+        x = self.fc(x)
+        # temoral average pooling
+        w = self.temporal_avg_pool(x, mask=x_mask)
+
+        return w

models/text_encoder.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn as nn
+
+from models.dit import DiTConVBlock
+
+def sequence_mask(length: torch.Tensor, max_length: int = None) -> torch.Tensor:
+    if max_length is None:
+        max_length = length.max()
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+# modified from https://github.com/jaywalnut310/vits/blob/main/models.py
+class TextEncoder(nn.Module):
+    def __init__(self, n_vocab, out_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, gin_channels):
+        super().__init__()
+        self.n_vocab = n_vocab
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        
+        self.scale = self.hidden_channels ** 0.5
+
+        self.emb = nn.Embedding(n_vocab, hidden_channels)
+        nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)
+
+        self.encoder = nn.ModuleList([DiTConVBlock(hidden_channels, filter_channels, n_heads, kernel_size, p_dropout, gin_channels) for _ in range(n_layers)])
+        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
+        
+        self.initialize_weights()
+        
+    def initialize_weights(self):
+        for block in self.encoder:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+    def forward(self, x: torch.Tensor, c: torch.Tensor, x_lengths: torch.Tensor):
+        x = self.emb(x) * self.scale  # [b, t, h]
+        x = x.transpose(1, -1)  # [b, h, t]
+        x_mask = sequence_mask(x_lengths, x.size(2)).unsqueeze(1).to(x.dtype)
+
+        for layer in self.encoder:
+            x = layer(x, c, x_mask)
+        mu_x = self.proj(x) * x_mask
+
+        return x, mu_x, x_mask

monotonic_align/__init__.py

@@ -0,0 +1,16 @@
+from numpy import zeros, int32, float32
+from torch import from_numpy
+
+from .core import maximum_path_jit
+
+
+def maximum_path(neg_cent, mask):
+    device = neg_cent.device
+    dtype = neg_cent.dtype
+    neg_cent = neg_cent.data.cpu().numpy().astype(float32)
+    path = zeros(neg_cent.shape, dtype=int32)
+
+    t_t_max = mask.sum(1)[:, 0].data.cpu().numpy().astype(int32)
+    t_s_max = mask.sum(2)[:, 0].data.cpu().numpy().astype(int32)
+    maximum_path_jit(path, neg_cent, t_t_max, t_s_max)
+    return from_numpy(path).to(device=device, dtype=dtype)

monotonic_align/core.py

@@ -0,0 +1,46 @@
+import numba
+
+
+@numba.jit(
+    numba.void(
+        numba.int32[:, :, ::1],
+        numba.float32[:, :, ::1],
+        numba.int32[::1],
+        numba.int32[::1],
+    ),
+    nopython=True,
+    nogil=True,
+)
+def maximum_path_jit(paths, values, t_ys, t_xs):
+    b = paths.shape[0]
+    max_neg_val = -1e9
+    for i in range(int(b)):
+        path = paths[i]
+        value = values[i]
+        t_y = t_ys[i]
+        t_x = t_xs[i]
+
+        v_prev = v_cur = 0.0
+        index = t_x - 1
+
+        for y in range(t_y):
+            for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)):
+                if x == y:
+                    v_cur = max_neg_val
+                else:
+                    v_cur = value[y - 1, x]
+                if x == 0:
+                    if y == 0:
+                        v_prev = 0.0
+                    else:
+                        v_prev = max_neg_val
+                else:
+                    v_prev = value[y - 1, x - 1]
+                value[y, x] += max(v_prev, v_cur)
+
+        for y in range(t_y - 1, -1, -1):
+            path[y, index] = 1
+            if index != 0 and (
+                index == y or value[y - 1, index] < value[y - 1, index - 1]
+            ):
+                index = index - 1

requirements.txt

@@ -0,0 +1,11 @@
+torch
+torchaudio
+matplotlib
+numpy
+tensorboard
+pypinyin
+jieba
+eng_to_ipa
+unidecode
+inflect
+pyopenjtalk-prebuilt

text/LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2017 Keith Ito
+
+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.

text/__init__.py

@@ -0,0 +1,71 @@
+""" from https://github.com/keithito/tacotron """
+from text import cleaners
+from text.symbols import symbols
+
+
+# Mappings from symbol to numeric ID and vice versa:
+_symbol_to_id = {s: i for i, s in enumerate(symbols)}
+_id_to_symbol = {i: s for i, s in enumerate(symbols)}
+
+
+def text_to_sequence(text, symbols, cleaner_names):
+  '''Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+    Args:
+      text: string to convert to a sequence
+      cleaner_names: names of the cleaner functions to run the text through
+    Returns:
+      List of integers corresponding to the symbols in the text
+  '''
+  sequence = []
+  symbol_to_id = {s: i for i, s in enumerate(symbols)}
+  clean_text = _clean_text(text, cleaner_names)
+  print(clean_text)
+  print(f" length:{len(clean_text)}")
+  for symbol in clean_text:
+    if symbol not in symbol_to_id.keys():
+      continue
+    symbol_id = symbol_to_id[symbol]
+    sequence += [symbol_id]
+  print(f" length:{len(sequence)}")
+  return sequence
+
+
+def cleaned_text_to_sequence(cleaned_text):
+  '''Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+    Args:
+      text: string to convert to a sequence
+    Returns:
+      List of integers corresponding to the symbols in the text
+  '''
+  # symbol_to_id = {s: i for i, s in enumerate(symbols)}
+  sequence = [_symbol_to_id[symbol] for symbol in cleaned_text if symbol in _symbol_to_id.keys()]
+  return sequence
+
+def cleaned_text_to_sequence_chinese(cleaned_text):
+  '''Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+    Args:
+      text: string to convert to a sequence
+    Returns:
+      List of integers corresponding to the symbols in the text
+  '''
+  # symbol_to_id = {s: i for i, s in enumerate(symbols)}
+  sequence = [_symbol_to_id[symbol] for symbol in cleaned_text.split(' ') if symbol in _symbol_to_id.keys()]
+  return sequence
+
+
+def sequence_to_text(sequence):
+  '''Converts a sequence of IDs back to a string'''
+  result = ''
+  for symbol_id in sequence:
+    s = _id_to_symbol[symbol_id]
+    result += s
+  return result
+
+
+def _clean_text(text, cleaner_names):
+  for name in cleaner_names:
+    cleaner = getattr(cleaners, name)
+    if not cleaner:
+      raise Exception('Unknown cleaner: %s' % name)
+    text = cleaner(text)
+  return text

text/cleaners.py

@@ -0,0 +1,73 @@
+import re
+import string
+import numpy as np
+from .langdetect import detect, LangDetectException
+
+from text.english import english_to_ipa2
+from text.mandarin import chinese_to_cnm3
+from text.japanese import japanese_to_ipa2
+
+language_module_map = {"PAD":0, "ZH": 1, "EN": 2, "JA": 3}
+
+# 预编译正则表达式
+ZH_PATTERN = re.compile(r'[\u3400-\u4DBF\u4e00-\u9FFF\uF900-\uFAFF\u3000-\u303F]')
+EN_PATTERN = re.compile(r'[a-zA-Z.,!?\'"(){}[\]<>:;@#$%^&*-_+=/\\|~`]+')
+JP_PATTERN = re.compile(r'[\u3040-\u309F\u30A0-\u30FF\u4E00-\u9FAF\u31F0-\u31FF\uFF00-\uFFEF\u3000-\u303F]')
+CLEANER_PATTERN = re.compile(r'\[(ZH|EN|JA)\]')
+
+def detect_language(text: str, prev_lang=None):
+    """
+    根据给定的文本检测语言
+
+    :param text: 输入文本
+    :param prev_lang: 上一个检测到的语言
+    :return: 'ZH' for Chinese, 'EN' for English, 'JA' for Japanese, or prev_lang for spaces
+    """
+    if ZH_PATTERN.search(text): return 'ZH'
+    if EN_PATTERN.search(text): return 'EN'
+    if JP_PATTERN.search(text): return 'JA'
+    if text.isspace(): return prev_lang  # 若是空格，则返回前一个语言
+    return None
+
+def replace_substring(s, start_index, end_index, replacement):
+    return s[:start_index] + replacement + s[end_index:]
+
+def replace_sublist(lst, start_index, end_index, replacement_list):
+    lst[start_index:end_index] = replacement_list
+
+# convert text to ipa and prepare for language embedding
+def append_tags_and_convert(match, conversion_func, tag_value, tags):
+    converted_text = conversion_func(match.group(1))
+    tags.extend([tag_value] * len(converted_text))
+    return converted_text + ' '
+
+# auto detect language using re
+def cjke_cleaners4(text: str):
+    """
+    根据文本内容自动检测语言并转换为IPA音标
+
+    :param text: 输入文本
+    :return: 转换为IPA音标的文本
+    """
+    text = CLEANER_PATTERN.sub('', text)
+    pointer = 0
+    output = ''
+    current_language = detect_language(text[pointer])
+    
+    while pointer < len(text):
+        temp_text = ''
+        while pointer < len(text) and detect_language(text[pointer], current_language) == current_language:
+            temp_text += text[pointer]
+            pointer += 1
+        if current_language == 'ZH':
+            output += chinese_to_cnm3(temp_text)
+        elif current_language == 'JA':
+            output += japanese_to_ipa2(temp_text)
+        elif current_language == 'EN':
+            output += english_to_ipa2(temp_text)
+        if pointer < len(text):
+            current_language = detect_language(text[pointer])
+
+    output = re.sub(r'\s+$', '', output)
+    output = re.sub(r'([^\.,!\?\-…~])$', r'\1.', output)
+    return output

text/cn2an/__init__.py

@@ -0,0 +1,16 @@
+__version__ = "0.5.20"
+
+from .cn2an import Cn2An
+from .an2cn import An2Cn
+from .transform import Transform
+
+cn2an = Cn2An().cn2an
+an2cn = An2Cn().an2cn
+transform = Transform().transform
+
+__all__ = [
+    "__version__",
+    "cn2an",
+    "an2cn",
+    "transform"
+]

text/cn2an/an2cn.py

@@ -0,0 +1,203 @@
+from typing import Union
+
+#from proces import preprocess
+
+from .conf import NUMBER_LOW_AN2CN, NUMBER_UP_AN2CN, UNIT_LOW_ORDER_AN2CN, UNIT_UP_ORDER_AN2CN
+
+
+class An2Cn(object):
+    def __init__(self) -> None:
+        self.all_num = "0123456789"
+        self.number_low = NUMBER_LOW_AN2CN
+        self.number_up = NUMBER_UP_AN2CN
+        self.mode_list = ["low", "up", "rmb", "direct"]
+
+    def an2cn(self, inputs: Union[str, int, float] = None, mode: str = "low") -> str:
+        """阿拉伯数字转中文数字
+
+        :param inputs: 阿拉伯数字
+        :param mode: low 小写数字，up 大写数字，rmb 人民币大写，direct 直接转化
+        :return: 中文数字
+        """
+        if inputs is not None and inputs != "":
+            if mode not in self.mode_list:
+                raise ValueError(f"mode 仅支持 {str(self.mode_list)} ！")
+
+            # 将数字转化为字符串，这里会有Python会自动做转化
+            # 1. -> 1.0 1.00 -> 1.0 -0 -> 0
+            if not isinstance(inputs, str):
+                inputs = self.__number_to_string(inputs)
+
+            # 数据预处理：
+            # 1. 繁体转简体
+            # 2. 全角转半角
+            # inputs = preprocess(inputs, pipelines=[
+            #     "traditional_to_simplified",
+            #     "full_angle_to_half_angle"
+            # ])
+
+            # 检查数据是否有效
+            self.__check_inputs_is_valid(inputs)
+
+            # 判断正负
+            if inputs[0] == "-":
+                sign = "负"
+                inputs = inputs[1:]
+            else:
+                sign = ""
+
+            if mode == "direct":
+                output = self.__direct_convert(inputs)
+            else:
+                # 切割整数部分和小数部分
+                split_result = inputs.split(".")
+                len_split_result = len(split_result)
+                if len_split_result == 1:
+                    # 不包含小数的输入
+                    integer_data = split_result[0]
+                    if mode == "rmb":
+                        output = self.__integer_convert(integer_data, "up") + "元整"
+                    else:
+                        output = self.__integer_convert(integer_data, mode)
+                elif len_split_result == 2:
+                    # 包含小数的输入
+                    integer_data, decimal_data = split_result
+                    if mode == "rmb":
+                        int_data = self.__integer_convert(integer_data, "up")
+                        dec_data = self.__decimal_convert(decimal_data, "up")
+                        len_dec_data = len(dec_data)
+
+                        if len_dec_data == 0:
+                            output = int_data + "元整"
+                        elif len_dec_data == 1:
+                            raise ValueError(f"异常输出：{dec_data}")
+                        elif len_dec_data == 2:
+                            if dec_data[1] != "零":
+                                if int_data == "零":
+                                    output = dec_data[1] + "角"
+                                else:
+                                    output = int_data + "元" + dec_data[1] + "角"
+                            else:
+                                output = int_data + "元整"
+                        else:
+                            if dec_data[1] != "零":
+                                if dec_data[2] != "零":
+                                    if int_data == "零":
+                                        output = dec_data[1] + "角" + dec_data[2] + "分"
+                                    else:
+                                        output = int_data + "元" + dec_data[1] + "角" + dec_data[2] + "分"
+                                else:
+                                    if int_data == "零":
+                                        output = dec_data[1] + "角"
+                                    else:
+                                        output = int_data + "元" + dec_data[1] + "角"
+                            else:
+                                if dec_data[2] != "零":
+                                    if int_data == "零":
+                                        output = dec_data[2] + "分"
+                                    else:
+                                        output = int_data + "元" + "零" + dec_data[2] + "分"
+                                else:
+                                    output = int_data + "元整"
+                    else:
+                        output = self.__integer_convert(integer_data, mode) + self.__decimal_convert(decimal_data, mode)
+                else:
+                    raise ValueError(f"输入格式错误：{inputs}！")
+        else:
+            raise ValueError("输入数据为空！")
+
+        return sign + output
+
+    def __direct_convert(self, inputs: str) -> str:
+        _output = ""
+        for d in inputs:
+            if d == ".":
+                _output += "点"
+            else:
+                _output += self.number_low[int(d)]
+        return _output
+
+    @staticmethod
+    def __number_to_string(number_data: Union[int, float]) -> str:
+        # 小数处理：python 会自动把 0.00005 转化成 5e-05，因此 str(0.00005) != "0.00005"
+        string_data = str(number_data)
+        if "e" in string_data:
+            string_data_list = string_data.split("e")
+            string_key = string_data_list[0]
+            string_value = string_data_list[1]
+            if string_value[0] == "-":
+                string_data = "0." + "0" * (int(string_value[1:]) - 1) + string_key
+            else:
+                string_data = string_key + "0" * int(string_value)
+        return string_data
+
+    def __check_inputs_is_valid(self, check_data: str) -> None:
+        # 检查输入数据是否在规定的字典中
+        all_check_keys = self.all_num + ".-"
+        for data in check_data:
+            if data not in all_check_keys:
+                raise ValueError(f"输入的数据不在转化范围内：{data}！")
+
+    def __integer_convert(self, integer_data: str, mode: str) -> str:
+        if mode == "low":
+            numeral_list = NUMBER_LOW_AN2CN
+            unit_list = UNIT_LOW_ORDER_AN2CN
+        elif mode == "up":
+            numeral_list = NUMBER_UP_AN2CN
+            unit_list = UNIT_UP_ORDER_AN2CN
+        else:
+            raise ValueError(f"error mode: {mode}")
+
+        # 去除前面的 0，比如 007 => 7
+        integer_data = str(int(integer_data))
+
+        len_integer_data = len(integer_data)
+        if len_integer_data > len(unit_list):
+            raise ValueError(f"超出数据范围，最长支持 {len(unit_list)} 位")
+
+        output_an = ""
+        for i, d in enumerate(integer_data):
+            if int(d):
+                output_an += numeral_list[int(d)] + unit_list[len_integer_data - i - 1]
+            else:
+                if not (len_integer_data - i - 1) % 4:
+                    output_an += numeral_list[int(d)] + unit_list[len_integer_data - i - 1]
+
+                if i > 0 and not output_an[-1] == "零":
+                    output_an += numeral_list[int(d)]
+
+        output_an = output_an.replace("零零", "零").replace("零万", "万").replace("零亿", "亿").replace("亿万", "亿") \
+            .strip("零")
+
+        # 解决「一十几」问题
+        if output_an[:2] in ["一十"]:
+            output_an = output_an[1:]
+
+        # 0 - 1 之间的小数
+        if not output_an:
+            output_an = "零"
+
+        return output_an
+
+    def __decimal_convert(self, decimal_data: str, o_mode: str) -> str:
+        len_decimal_data = len(decimal_data)
+
+        if len_decimal_data > 16:
+            print(f"注意：小数部分长度为 {len_decimal_data} ，将自动截取前 16 位有效精度！")
+            decimal_data = decimal_data[:16]
+
+        if len_decimal_data:
+            output_an = "点"
+        else:
+            output_an = ""
+
+        if o_mode == "low":
+            numeral_list = NUMBER_LOW_AN2CN
+        elif o_mode == "up":
+            numeral_list = NUMBER_UP_AN2CN
+        else:
+            raise ValueError(f"error mode: {o_mode}")
+
+        for data in decimal_data:
+            output_an += numeral_list[int(data)]
+        return output_an

text/cn2an/cn2an.py

@@ -0,0 +1,293 @@
+import re
+from typing import Union
+
+#from proces import preprocess
+
+from .an2cn import An2Cn
+from .conf import NUMBER_CN2AN, UNIT_CN2AN, STRICT_CN_NUMBER, NORMAL_CN_NUMBER, NUMBER_LOW_AN2CN, UNIT_LOW_AN2CN
+
+
+class Cn2An(object):
+    def __init__(self) -> None:
+        self.all_num = "".join(list(NUMBER_CN2AN.keys()))
+        self.all_unit = "".join(list(UNIT_CN2AN.keys()))
+        self.strict_cn_number = STRICT_CN_NUMBER
+        self.normal_cn_number = NORMAL_CN_NUMBER
+        self.check_key_dict = {
+            "strict": "".join(self.strict_cn_number.values()) + "点负",
+            "normal": "".join(self.normal_cn_number.values()) + "点负",
+            "smart": "".join(self.normal_cn_number.values()) + "点负" + "01234567890.-"
+        }
+        self.pattern_dict = self.__get_pattern()
+        self.ac = An2Cn()
+        self.mode_list = ["strict", "normal", "smart"]
+        self.yjf_pattern = re.compile(fr"^.*?[元圆][{self.all_num}]角([{self.all_num}]分)?$")
+        self.pattern1 = re.compile(fr"^-?\d+(\.\d+)?[{self.all_unit}]?$")
+        self.ptn_all_num = re.compile(f"^[{self.all_num}]+$")
+        # "十?" is for special case "十一万三"
+        self.ptn_speaking_mode = re.compile(f"^([{self.all_num}]{{0,2}}[{self.all_unit}])+[{self.all_num}]$")
+
+    def cn2an(self, inputs: Union[str, int, float] = None, mode: str = "strict") -> Union[float, int]:
+        """中文数字转阿拉伯数字
+
+        :param inputs: 中文数字、阿拉伯数字、中文数字和阿拉伯数字
+        :param mode: strict 严格，normal 正常，smart 智能
+        :return: 阿拉伯数字
+        """
+        if inputs is not None or inputs == "":
+            if mode not in self.mode_list:
+                raise ValueError(f"mode 仅支持 {str(self.mode_list)} ！")
+
+            # 将数字转化为字符串
+            if not isinstance(inputs, str):
+                inputs = str(inputs)
+
+            # 数据预处理：
+            # 1. 繁体转简体
+            # 2. 全角转半角
+            # inputs = preprocess(inputs, pipelines=[
+            #     "traditional_to_simplified",
+            #     "full_angle_to_half_angle"
+            # ])
+
+            # 特殊转化 廿
+            inputs = inputs.replace("廿", "二十")
+
+            # 检查输入数据是否有效
+            sign, integer_data, decimal_data, is_all_num = self.__check_input_data_is_valid(inputs, mode)
+
+            # smart 下的特殊情况
+            if sign == 0:
+                return integer_data
+            else:
+                if not is_all_num:
+                    if decimal_data is None:
+                        output = self.__integer_convert(integer_data)
+                    else:
+                        output = self.__integer_convert(integer_data) + self.__decimal_convert(decimal_data)
+                        # fix 1 + 0.57 = 1.5699999999999998
+                        output = round(output, len(decimal_data))
+                else:
+                    if decimal_data is None:
+                        output = self.__direct_convert(integer_data)
+                    else:
+                        output = self.__direct_convert(integer_data) + self.__decimal_convert(decimal_data)
+                        # fix 1 + 0.57 = 1.5699999999999998
+                        output = round(output, len(decimal_data))
+        else:
+            raise ValueError("输入数据为空！")
+
+        return sign * output
+
+    def __get_pattern(self) -> dict:
+        # 整数严格检查
+        _0 = "[零]"
+        _1_9 = "[一二三四五六七八九]"
+        _10_99 = f"{_1_9}?[十]{_1_9}?"
+        _1_99 = f"({_10_99}|{_1_9})"
+        _100_999 = f"({_1_9}[百]([零]{_1_9})?|{_1_9}[百]{_10_99})"
+        _1_999 = f"({_100_999}|{_1_99})"
+        _1000_9999 = f"({_1_9}[千]([零]{_1_99})?|{_1_9}[千]{_100_999})"
+        _1_9999 = f"({_1000_9999}|{_1_999})"
+        _10000_99999999 = f"({_1_9999}[万]([零]{_1_999})?|{_1_9999}[万]{_1000_9999})"
+        _1_99999999 = f"({_10000_99999999}|{_1_9999})"
+        _100000000_9999999999999999 = f"({_1_99999999}[亿]([零]{_1_99999999})?|{_1_99999999}[亿]{_10000_99999999})"
+        _1_9999999999999999 = f"({_100000000_9999999999999999}|{_1_99999999})"
+        str_int_pattern = f"^({_0}|{_1_9999999999999999})$"
+        nor_int_pattern = f"^({_0}|{_1_9999999999999999})$"
+
+        str_dec_pattern = "^[零一二三四五六七八九]{0,15}[一二三四五六七八九]$"
+        nor_dec_pattern = "^[零一二三四五六七八九]{0,16}$"
+
+        for str_num in self.strict_cn_number.keys():
+            str_int_pattern = str_int_pattern.replace(str_num, self.strict_cn_number[str_num])
+            str_dec_pattern = str_dec_pattern.replace(str_num, self.strict_cn_number[str_num])
+        for nor_num in self.normal_cn_number.keys():
+            nor_int_pattern = nor_int_pattern.replace(nor_num, self.normal_cn_number[nor_num])
+            nor_dec_pattern = nor_dec_pattern.replace(nor_num, self.normal_cn_number[nor_num])
+
+        pattern_dict = {
+            "strict": {
+                "int": re.compile(str_int_pattern),
+                "dec": re.compile(str_dec_pattern)
+            },
+            "normal": {
+                "int": re.compile(nor_int_pattern),
+                "dec": re.compile(nor_dec_pattern)
+            }
+        }
+        return pattern_dict
+
+    def __copy_num(self, num):
+        cn_num = ""
+        for n in num:
+            cn_num += NUMBER_LOW_AN2CN[int(n)]
+        return cn_num
+
+    def __check_input_data_is_valid(self, check_data: str, mode: str) -> (int, str, str, bool):
+        # 去除 元整、圆整、元正、圆正
+        stop_words = ["元整", "圆整", "元正", "圆正"]
+        for word in stop_words:
+            if check_data[-2:] == word:
+                check_data = check_data[:-2]
+
+        # 去除 元、圆
+        if mode != "strict":
+            normal_stop_words = ["圆", "元"]
+            for word in normal_stop_words:
+                if check_data[-1] == word:
+                    check_data = check_data[:-1]
+
+        # 处理元角分
+        result = self.yjf_pattern.search(check_data)
+        if result:
+            check_data = check_data.replace("元", "点").replace("角", "").replace("分", "")
+
+        # 处理特殊问法：一千零十一 一万零百一十一
+        if "零十" in check_data:
+            check_data = check_data.replace("零十", "零一十")
+        if "零百" in check_data:
+            check_data = check_data.replace("零百", "零一百")
+
+        for data in check_data:
+            if data not in self.check_key_dict[mode]:
+                raise ValueError(f"当前为{mode}模式，输入的数据不在转化范围内：{data}！")
+
+        # 确定正负号
+        if check_data[0] == "负":
+            check_data = check_data[1:]
+            sign = -1
+        else:
+            sign = 1
+
+        if "点" in check_data:
+            split_data = check_data.split("点")
+            if len(split_data) == 2:
+                integer_data, decimal_data = split_data
+                # 将 smart 模式中的阿拉伯数字转化成中文数字
+                if mode == "smart":
+                    integer_data = re.sub(r"\d+", lambda x: self.ac.an2cn(x.group()), integer_data)
+                    decimal_data = re.sub(r"\d+", lambda x: self.__copy_num(x.group()), decimal_data)
+                    mode = "normal"
+            else:
+                raise ValueError("数据中包含不止一个点！")
+        else:
+            integer_data = check_data
+            decimal_data = None
+            # 将 smart 模式中的阿拉伯数字转化成中文数字
+            if mode == "smart":
+                # 10.1万 10.1
+                result1 = self.pattern1.search(integer_data)
+                if result1:
+                    if result1.group() == integer_data:
+                        if integer_data[-1] in UNIT_CN2AN.keys():
+                            output = int(float(integer_data[:-1]) * UNIT_CN2AN[integer_data[-1]])
+                        else:
+                            output = float(integer_data)
+                        return 0, output, None, None
+
+                integer_data = re.sub(r"\d+", lambda x: self.ac.an2cn(x.group()), integer_data)
+                mode = "normal"
+
+        result_int = self.pattern_dict[mode]["int"].search(integer_data)
+        if result_int:
+            if result_int.group() == integer_data:
+                if decimal_data is not None:
+                    result_dec = self.pattern_dict[mode]["dec"].search(decimal_data)
+                    if result_dec:
+                        if result_dec.group() == decimal_data:
+                            return sign, integer_data, decimal_data, False
+                else:
+                    return sign, integer_data, decimal_data, False
+        else:
+            if mode == "strict":
+                raise ValueError(f"不符合格式的数据：{integer_data}")
+            elif mode == "normal":
+                # 纯数模式：一二三
+                result_all_num = self.ptn_all_num.search(integer_data)
+                if result_all_num:
+                    if result_all_num.group() == integer_data:
+                        if decimal_data is not None:
+                            result_dec = self.pattern_dict[mode]["dec"].search(decimal_data)
+                            if result_dec:
+                                if result_dec.group() == decimal_data:
+                                    return sign, integer_data, decimal_data, True
+                        else:
+                            return sign, integer_data, decimal_data, True
+
+                # 口语模式：一万二，两千三，三百四，十三万六，一百二十五万三
+                result_speaking_mode = self.ptn_speaking_mode.search(integer_data)
+                if len(integer_data) >= 3 and result_speaking_mode and result_speaking_mode.group() == integer_data:
+                    # len(integer_data)>=3: because the minimum length of integer_data that can be matched is 3
+                    # to find the last unit
+                    last_unit = result_speaking_mode.groups()[-1][-1]
+                    _unit = UNIT_LOW_AN2CN[UNIT_CN2AN[last_unit] // 10]
+                    integer_data = integer_data + _unit
+                    if decimal_data is not None:
+                        result_dec = self.pattern_dict[mode]["dec"].search(decimal_data)
+                        if result_dec:
+                            if result_dec.group() == decimal_data:
+                                return sign, integer_data, decimal_data, False
+                    else:
+                        return sign, integer_data, decimal_data, False
+
+        raise ValueError(f"不符合格式的数据：{check_data}")
+
+    def __integer_convert(self, integer_data: str) -> int:
+        # 核心
+        output_integer = 0
+        unit = 1
+        ten_thousand_unit = 1
+        for index, cn_num in enumerate(reversed(integer_data)):
+            # 数值
+            if cn_num in NUMBER_CN2AN:
+                num = NUMBER_CN2AN[cn_num]
+                output_integer += num * unit
+            # 单位
+            elif cn_num in UNIT_CN2AN:
+                unit = UNIT_CN2AN[cn_num]
+                # 判断出万、亿、万亿
+                if unit % 10000 == 0:
+                    # 万 亿
+                    if unit > ten_thousand_unit:
+                        ten_thousand_unit = unit
+                    # 万亿
+                    else:
+                        ten_thousand_unit = unit * ten_thousand_unit
+                        unit = ten_thousand_unit
+
+                if unit < ten_thousand_unit:
+                    unit = unit * ten_thousand_unit
+
+                if index == len(integer_data) - 1:
+                    output_integer += unit
+            else:
+                raise ValueError(f"{cn_num} 不在转化范围内")
+
+        return int(output_integer)
+
+    def __decimal_convert(self, decimal_data: str) -> float:
+        len_decimal_data = len(decimal_data)
+
+        if len_decimal_data > 16:
+            print(f"注意：小数部分长度为 {len_decimal_data} ，将自动截取前 16 位有效精度！")
+            decimal_data = decimal_data[:16]
+            len_decimal_data = 16
+
+        output_decimal = 0
+        for index in range(len(decimal_data) - 1, -1, -1):
+            unit_key = NUMBER_CN2AN[decimal_data[index]]
+            output_decimal += unit_key * 10 ** -(index + 1)
+
+        # 处理精度溢出问题
+        output_decimal = round(output_decimal, len_decimal_data)
+
+        return output_decimal
+
+    def __direct_convert(self, data: str) -> int:
+        output_data = 0
+        for index in range(len(data) - 1, -1, -1):
+            unit_key = NUMBER_CN2AN[data[index]]
+            output_data += unit_key * 10 ** (len(data) - index - 1)
+
+        return output_data