Upload 238 files

.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.english import english_to_ipa2
+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 = '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 = english_to_ipa2
+    
+    # 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=0.667, 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')) + list(Path('./audios').rglob('*.flac'))
+
+    # 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="Today I want to tell you three stories from my life. That's it. No big deal. Just three stories.",
+                )
+             
+                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=40,
+                    value=8,
+                    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:7c473c5dc40bf87e8472f8688790cb20a2ec10494fa08cb710d657cf1f892d44
+size 37552627

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,12 @@
+torch
+torchaudio
+matplotlib
+numpy
+tensorboard
+pypinyin
+jieba
+eng_to_ipa
+unidecode
+inflect
+pyopenjtalk-prebuilt
+numba

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