hfdemo

README.md

@@ -1,13 +1,13 @@
 ---
-title: Pits
+title: PITS
 emoji: 🌖
 colorFrom: purple
 colorTo: indigo
 sdk: gradio
-sdk_version: 3.20.0
+sdk_version: 3.20.1
 app_file: app.py
-pinned: false
+pinned: true
 license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Official Demo for PITS: Variational Pitch Inference without Fundamental Frequency for End-to-End Pitch-controllable TTS

analysis.py

@@ -0,0 +1,141 @@
+# modified from https://github.com/dhchoi99/NANSY
+# We have modified the implementation of dhchoi99 to be fully differentiable.
+import math
+import torch
+from yin import *
+
+
+class Pitch(torch.nn.Module):
+
+    def __init__(
+            self, 
+            sr=22050, 
+            w_step=256, 
+            W=2048, 
+            tau_max=2048, 
+            midi_start=5, 
+            midi_end=85, 
+            octave_range=12):
+        super(Pitch, self).__init__()
+        self.sr = sr
+        self.w_step = w_step
+        self.W = W
+        self.tau_max = tau_max
+        self.unfold = torch.nn.Unfold((1, self.W),
+                                      1,
+                                      0,
+                                      stride=(1, self.w_step))
+        midis = list(range(midi_start, midi_end))
+        self.len_midis = len(midis)
+        c_ms = torch.tensor([self.midi_to_lag(m, octave_range) for m in midis])
+        self.register_buffer('c_ms', c_ms)
+        self.register_buffer('c_ms_ceil', torch.ceil(self.c_ms).long())
+        self.register_buffer('c_ms_floor', torch.floor(self.c_ms).long())
+
+    def midi_to_lag(self, m: int, octave_range: float = 12):
+        """converts midi-to-lag, eq. (4)
+
+        Args:
+            m: midi
+            sr: sample_rate
+            octave_range:
+
+        Returns:
+            lag: time lag(tau, c(m)) calculated from midi, eq. (4)
+
+        """
+        f = 440 * math.pow(2, (m - 69) / octave_range)
+        lag = self.sr / f
+        return lag
+
+    def yingram_from_cmndf(self, cmndfs: torch.Tensor) -> torch.Tensor:
+        """ yingram calculator from cMNDFs(cumulative Mean Normalized Difference Functions)
+
+        Args:
+            cmndfs: torch.Tensor
+                calculated cumulative mean normalized difference function
+                for details, see models/yin.py or eq. (1) and (2)
+            ms: list of midi(int)
+            sr: sampling rate
+
+        Returns:
+            y:
+                calculated batch yingram
+
+
+        """
+        #c_ms = np.asarray([Pitch.midi_to_lag(m, sr) for m in ms])
+        #c_ms = torch.from_numpy(c_ms).to(cmndfs.device)
+
+        y = (cmndfs[:, self.c_ms_ceil] -
+             cmndfs[:, self.c_ms_floor]) / (self.c_ms_ceil - self.c_ms_floor).unsqueeze(0) * (
+                 self.c_ms - self.c_ms_floor).unsqueeze(0) + cmndfs[:, self.c_ms_floor]
+        return y
+
+    def yingram(self, x: torch.Tensor):
+        """calculates yingram from raw audio (multi segment)
+
+        Args:
+            x: raw audio, torch.Tensor of shape (t)
+            W: yingram Window Size
+            tau_max:
+            sr: sampling rate
+            w_step: yingram bin step size
+
+        Returns:
+            yingram: yingram. torch.Tensor of shape (80 x t')
+
+        """
+        # x.shape: t -> B,T, B,T = x.shape
+        B, T = x.shape
+        w_len = self.W
+
+
+        frames = self.unfold(x.view(B, 1, 1, T))
+        frames = frames.permute(0, 2,
+                                1).contiguous().view(-1,
+                                                     self.W)  #[B* frames, W]
+        # If not using gpu, or torch not compatible, implemented numpy batch function is still fine
+        dfs = differenceFunctionTorch(frames, frames.shape[-1], self.tau_max)
+        cmndfs = cumulativeMeanNormalizedDifferenceFunctionTorch(
+            dfs, self.tau_max)
+        yingram = self.yingram_from_cmndf(cmndfs)  #[B*frames,F]
+        yingram = yingram.view(B, -1, self.len_midis).permute(0, 2,
+                                                              1)  # [B,F,T]
+        return yingram
+
+    def crop_scope(self, x, yin_start,
+                   scope_shift):  # x: tensor [B,C,T] #scope_shift: tensor [B]
+        return torch.stack([
+            x[i, yin_start + scope_shift[i]:yin_start + self.yin_scope +
+              scope_shift[i], :] for i in range(x.shape[0])
+        ],
+                           dim=0)
+
+
+if __name__ == '__main__':
+    import torch
+    import librosa as rosa
+    import matplotlib.pyplot as plt
+    wav = torch.tensor(rosa.load('LJ001-0002.wav', sr=22050,
+                                 mono=True)[0]).unsqueeze(0)
+    #    wav = torch.randn(1,40965)
+
+    wav = torch.nn.functional.pad(wav, (0, (-wav.shape[1]) % 256))
+    #    wav = wav[#:,:8096]
+    print(wav.shape)
+    pitch = Pitch()
+
+    with torch.no_grad():
+        ps = pitch.yingram(torch.nn.functional.pad(wav, (1024, 1024)))
+        ps = torch.nn.functional.pad(ps, (0, 0, 8, 8), mode='replicate')
+        print(ps.shape)
+        spec = torch.stft(wav, 1024, 256, return_complex=False)
+        print(spec.shape)
+        plt.subplot(2, 1, 1)
+        plt.pcolor(ps[0].numpy(), cmap='magma')
+        plt.colorbar()
+        plt.subplot(2, 1, 2)
+        plt.pcolor(ps[0][15:65, :].numpy(), cmap='magma')
+        plt.colorbar()
+        plt.show()

app.py

@@ -0,0 +1,141 @@
+import gradio as gr
+import argparse
+import torch
+import commons
+import utils
+from models import (
+    SynthesizerTrn, )
+
+from text.symbols import symbol_len, lang_to_dict
+
+# we use Kyubyong/g2p for demo instead of our internal g2p
+# https://github.com/Kyubyong/g2p
+from g2p_en import G2p
+import re
+
+_symbol_to_id = lang_to_dict("en_US")
+
+class GradioApp:
+
+    def __init__(self, args):
+        self.hps = utils.get_hparams_from_file(args.config)
+        self.device = "cpu"
+        self.net_g = SynthesizerTrn(symbol_len(self.hps.data.languages),
+                                    self.hps.data.filter_length // 2 + 1,
+                                    self.hps.train.segment_size //
+                                    self.hps.data.hop_length,
+                                    midi_start=-5,
+                                    midi_end=75,
+                                    octave_range=24,
+                                    n_speakers=len(self.hps.data.speakers),
+                                    **self.hps.model).to(self.device)
+        _ = self.net_g.eval()
+        _ = utils.load_checkpoint(args.checkpoint_path, model_g=self.net_g)
+        self.g2p = G2p()
+        self.interface = self._gradio_interface()
+
+    def get_phoneme(self, text):
+        phones = [re.sub("[0-9]", "", p) for p in self.g2p(text)]
+        tone = [0 for p in phones]
+        if self.hps.data.add_blank:
+            text_norm = [_symbol_to_id[symbol] for symbol in phones]
+            text_norm = commons.intersperse(text_norm, 0)
+            tone = commons.intersperse(tone, 0)
+        else:
+            text_norm = phones
+        text_norm = torch.LongTensor(text_norm)
+        tone = torch.LongTensor(tone)
+        return text_norm, tone, phones
+    
+    @torch.no_grad()
+    def inference(self, text, speaker_id_val, seed, scope_shift, duration):
+        torch.manual_seed(seed)
+        text_norm, tone, phones = self.get_phoneme(text)
+        x_tst = text_norm.to(self.device).unsqueeze(0)
+        t_tst = tone.to(self.device).unsqueeze(0)
+        x_tst_lengths = torch.LongTensor([text_norm.size(0)]).to(self.device)
+        speaker_id = torch.LongTensor([speaker_id_val]).to(self.device)
+        decoder_inputs,*_ = self.net_g.infer_pre_decoder(
+                                           x_tst,
+                                           t_tst,
+                                           x_tst_lengths,
+                                           sid=speaker_id,
+                                           noise_scale=0.667,
+                                           noise_scale_w=0.8,
+                                           length_scale=duration,
+                                           scope_shift=scope_shift)
+        audio = self.net_g.infer_decode_chunk(
+            decoder_inputs, sid=speaker_id)[0, 0].data.cpu().float().numpy()
+        del decoder_inputs,  
+        return phones, (self.hps.data.sampling_rate, audio)
+
+
+    def _gradio_interface(self):
+        title = "PITS Demo"
+        inputs = [
+            gr.Textbox(label="Text (150 words limitation)",
+                       value="This is demo page.",
+                       elem_id="tts-input"),
+            gr.Dropdown(list(self.hps.data.speakers),
+                        value="p225",
+                        label="Speaker Identity",
+                        type="index"),
+            gr.Slider(0, 65536, step=1, label="random seed"),
+            gr.Slider(-15, 15, value=0, step=1, label="scope-shift"),
+            gr.Slider(0.5, 2., value=1., step=0.1,
+                      label="duration multiplier"),
+        ]
+        outputs = [
+            gr.Textbox(label="Phonemes"),
+            gr.Audio(type="numpy", label="Output audio")
+        ]
+        description = "Welcome to the Gradio demo for PITS: Variational Pitch Inference without Fundamental Frequency for End-to-End Pitch-controllable TTS.\n In this demo, we utilize an open-source G2P library (g2p_en) with stress removing, instead of our internal G2P.\n You can fix the latent z by controlling random seed.\n You can shift the pitch scope, but please note that this is opposite to pitch-shift. In addition, it is cropped from fixed z so please check pitch-controllability by comparing with normal synthesis.\n Thank you for trying out our PITS demo!"
+        article = "Github:https://github.com/anonymous-pits/pits \n Our current preprint contains several errors. Please wait for next update."
+        examples = [["This is a demo page of the PITS."],["I love hugging face."]]
+        return gr.Interface(
+            fn=self.inference,
+            inputs=inputs,
+            outputs=outputs,
+            title=title,
+            description=description,
+            article=article,
+            examples=examples,
+        )
+
+    def launch(self):
+        return self.interface.launch(share=True)
+
+
+def parsearg():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-c',
+                        '--config',
+                        type=str,
+                        default="./configs/config_en.yaml",
+                        help='Path to configuration file')
+    parser.add_argument('-m',
+                        '--model',
+                        type=str,
+                        default='PITS',
+                        help='Model name')
+    parser.add_argument('-r',
+                        '--checkpoint_path',
+                        type=str,
+                        default='./logs/pits_vctk_AD_3000.pth',
+                        help='Path to checkpoint for resume')
+    parser.add_argument('-f',
+                        '--force_resume',
+                        type=str,
+                        help='Path to checkpoint for force resume')
+    parser.add_argument('-d',
+                        '--dir',
+                        type=str,
+                        default='/DATA/audio/pits_samples',
+                        help='root dir')
+    args = parser.parse_args()
+    return args
+
+if __name__ == "__main__":
+    args = parsearg()
+    app = GradioApp(args)
+    app.launch()

attentions.py

@@ -0,0 +1,472 @@
+# from https://github.com/jaywalnut310/vits
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+import commons
+from modules import LayerNorm
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self, 
+        hidden_channels, 
+        filter_channels, 
+        n_heads, 
+        n_layers, 
+        kernel_size=1, 
+        p_dropout=0., 
+        window_size=4, 
+        **kwargs
+    ):
+        super().__init__()
+        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.window_size = window_size
+
+        self.drop = nn.Dropout(p_dropout)
+        self.attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+        for i in range(self.n_layers):
+            self.attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels, 
+                    hidden_channels, 
+                    n_heads, 
+                    p_dropout=p_dropout, 
+                    window_size=window_size
+                )
+            )
+            self.norm_layers_1.append(LayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FFN(
+                    hidden_channels, 
+                    hidden_channels, 
+                    filter_channels, 
+                    kernel_size, 
+                    p_dropout=p_dropout
+                )
+            )
+            self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+    def forward(self, x, x_mask):
+        attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            y = self.attn_layers[i](x, x, attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self, 
+        hidden_channels, 
+        filter_channels, 
+        n_heads, 
+        n_layers, 
+        kernel_size=1, 
+        p_dropout=0., 
+        proximal_bias=False, 
+        proximal_init=True, 
+        **kwargs
+    ):
+        super().__init__()
+        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.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+
+        self.drop = nn.Dropout(p_dropout)
+        self.self_attn_layers = nn.ModuleList()
+        self.norm_layers_0 = nn.ModuleList()
+        self.encdec_attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+        for i in range(self.n_layers):
+            self.self_attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels, 
+                    hidden_channels, 
+                    n_heads,
+                    p_dropout=p_dropout, 
+                    proximal_bias=proximal_bias, 
+                    proximal_init=proximal_init
+                )
+            )
+            self.norm_layers_0.append(LayerNorm(hidden_channels))
+            self.encdec_attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels, 
+                    hidden_channels, 
+                    n_heads, 
+                    p_dropout=p_dropout
+                )
+            )
+            self.norm_layers_1.append(LayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FFN(
+                    hidden_channels, 
+                    hidden_channels,
+                    filter_channels, 
+                    kernel_size, 
+                    p_dropout=p_dropout, 
+                    causal=True
+                )
+            )
+            self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+    def forward(self, x, x_mask, h, h_mask):
+        """
+        x: decoder input
+        h: encoder output
+        """
+        self_attn_mask = commons.subsequent_mask(
+            x_mask.size(2)
+        ).to(device=x.device, dtype=x.dtype)
+        encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            y = self.self_attn_layers[i](x, x, self_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_0[i](x + y)
+
+            y = self.encdec_attn_layers[i](x, h, encdec_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self, 
+        channels, 
+        out_channels, 
+        n_heads, 
+        p_dropout=0., 
+        window_size=None, 
+        heads_share=True, 
+        block_length=None, 
+        proximal_bias=False, 
+        proximal_init=False
+    ):
+        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.window_size = window_size
+        self.heads_share = heads_share
+        self.block_length = block_length
+        self.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+        self.attn = None
+
+        self.k_channels = channels // n_heads
+        self.conv_q = nn.Conv1d(channels, channels, 1)
+        self.conv_k = nn.Conv1d(channels, channels, 1)
+        self.conv_v = nn.Conv1d(channels, channels, 1)
+        self.conv_o = nn.Conv1d(channels, out_channels, 1)
+        self.drop = nn.Dropout(p_dropout)
+
+        if window_size is not None:
+            n_heads_rel = 1 if heads_share else n_heads
+            rel_stddev = self.k_channels**-0.5
+            self.emb_rel_k = nn.Parameter(torch.randn(
+                n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev)
+            self.emb_rel_v = nn.Parameter(torch.randn(
+                n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev)
+
+        nn.init.xavier_uniform_(self.conv_q.weight)
+        nn.init.xavier_uniform_(self.conv_k.weight)
+        nn.init.xavier_uniform_(self.conv_v.weight)
+        if proximal_init:
+            with torch.no_grad():
+                self.conv_k.weight.copy_(self.conv_q.weight)
+                self.conv_k.bias.copy_(self.conv_q.bias)
+
+    def forward(self, x, c, attn_mask=None):
+        q = self.conv_q(x)
+        k = self.conv_k(c)
+        v = self.conv_v(c)
+
+        x, self.attn = self.attention(q, k, v, mask=attn_mask)
+
+        x = self.conv_o(x)
+        return x
+
+    def attention(self, query, key, value, mask=None):
+        # reshape [b, d, t] -> [b, n_h, t, d_k]
+        b, d, t_s, t_t = (*key.size(), query.size(2))
+        #query = query.view(
+        #    b, 
+        #    self.n_heads, 
+        #    self.k_channels,
+        #    t_t
+        #).transpose(2, 3) #[b,h,t_t,c], d=h*c
+        #key = key.view(
+        #    b, 
+        #    self.n_heads, 
+        #    self.k_channels, 
+        #    t_s
+        #).transpose(2, 3) #[b,h,t_s,c]
+        #value = value.view(
+        #    b, 
+        #    self.n_heads, 
+        #    self.k_channels,
+        #    t_s
+        #).transpose(2, 3) #[b,h,t_s,c]
+        #scores = torch.matmul(
+        #    query / math.sqrt(self.k_channels), key.transpose(-2, -1)
+        #) #[b,h,t_t,t_s]
+        query = query.view(
+            b, 
+            self.n_heads, 
+            self.k_channels,
+            t_t
+        ) #[b,h,c,t_t]
+        key = key.view(
+            b, 
+            self.n_heads, 
+            self.k_channels, 
+            t_s
+        ) #[b,h,c,t_s]
+        value = value.view(
+            b, 
+            self.n_heads, 
+            self.k_channels,
+            t_s
+        ) #[b,h,c,t_s]
+        scores = torch.einsum('bhdt,bhds -> bhts', query / math.sqrt(self.k_channels), key) #[b,h,t_t,t_s]
+        #if self.window_size is not None:
+        #    assert t_s == t_t, "Relative attention is only available for self-attention."
+        #    key_relative_embeddings = self._get_relative_embeddings(
+        #        self.emb_rel_k, t_s
+        #    )
+        #    rel_logits = self._matmul_with_relative_keys(
+        #        query / math.sqrt(self.k_channels), key_relative_embeddings
+        #    ) #[b,h,t_t,d],[h or 1,e,d] ->[b,h,t_t,e]
+        #    scores_local = self._relative_position_to_absolute_position(rel_logits)
+        #    scores = scores + scores_local
+        #if self.proximal_bias:
+        #    assert t_s == t_t, "Proximal bias is only available for self-attention."
+        #    scores = scores + \
+        #        self._attention_bias_proximal(t_s).to(device=scores.device, dtype=scores.dtype)
+        #if mask is not None:
+        #    scores = scores.masked_fill(mask == 0, -1e4)
+        #    if self.block_length is not None:
+        #        assert t_s == t_t, "Local attention is only available for self-attention."
+        #        block_mask = torch.ones_like(scores).triu(-self.block_length).tril(self.block_length)
+        #        scores = scores.masked_fill(block_mask == 0, -1e4)
+        #p_attn = F.softmax(scores, dim=-1)  # [b, h, t_t, t_s]
+        #p_attn = self.drop(p_attn)
+        #output = torch.matmul(p_attn, value) # [b,h,t_t,t_s],[b,h,t_s,c] -> [b,h,t_t,c]
+        #if self.window_size is not None:
+        #    relative_weights = self._absolute_position_to_relative_position(p_attn) #[b, h, t_t, 2*t_t-1]
+        #    value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, t_s) #[h or 1, 2*t_t-1, c]
+        #    output = output + \
+        #        self._matmul_with_relative_values(
+        #            relative_weights, value_relative_embeddings) # [b, h, t_t, 2*t_t-1],[h or 1, 2*t_t-1, c] -> [b, h, t_t, c]
+        #output = output.transpose(2, 3).contiguous().view(b, d, t_t)  # [b, n_h, t_t, c] -> [b,h,c,t_t] -> [b, d, t_t]
+        if self.window_size is not None:
+            assert t_s == t_t, "Relative attention is only available for self-attention."
+            key_relative_embeddings = self._get_relative_embeddings(
+                self.emb_rel_k, t_s
+            )
+            rel_logits = torch.einsum('bhdt,hed->bhte',
+                query / math.sqrt(self.k_channels), key_relative_embeddings
+            ) #[b,h,c,t_t],[h or 1,e,c] ->[b,h,t_t,e]
+            scores_local = self._relative_position_to_absolute_position(rel_logits)
+            scores = scores + scores_local
+        if self.proximal_bias:
+            assert t_s == t_t, "Proximal bias is only available for self-attention."
+            scores = scores + \
+                self._attention_bias_proximal(t_s).to(device=scores.device, dtype=scores.dtype)
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e4)
+            if self.block_length is not None:
+                assert t_s == t_t, "Local attention is only available for self-attention."
+                block_mask = torch.ones_like(scores).triu(-self.block_length).tril(self.block_length)
+                scores = scores.masked_fill(block_mask == 0, -1e4)
+        p_attn = F.softmax(scores, dim=-1)  # [b, h, t_t, t_s]
+        p_attn = self.drop(p_attn)
+        output = torch.einsum('bhcs,bhts->bhct', value , p_attn) # [b,h,c,t_s],[b,h,t_t,t_s] -> [b,h,c,t_t]
+        if self.window_size is not None:
+            relative_weights = self._absolute_position_to_relative_position(p_attn) #[b, h, t_t, 2*t_t-1]
+            value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, t_s) #[h or 1, 2*t_t-1, c]
+            output = output + \
+                torch.einsum('bhte,hec->bhct',
+                    relative_weights, value_relative_embeddings) # [b, h, t_t, 2*t_t-1],[h or 1, 2*t_t-1, c] -> [b, h, c, t_t]
+        output = output.view(b, d, t_t)  # [b, h, c, t_t] -> [b, d, t_t]
+        return output, p_attn
+
+    def _matmul_with_relative_values(self, x, y):
+        """
+        x: [b, h, l, m]
+        y: [h or 1, m, d]
+        ret: [b, h, l, d]
+        """
+        ret = torch.matmul(x, y.unsqueeze(0))
+        return ret
+
+    def _matmul_with_relative_keys(self, x, y):
+        """
+        x: [b, h, l, d]
+        y: [h or 1, m, d]
+        ret: [b, h, l, m]
+        """
+        #ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
+        ret = torch.einsum('bhld,hmd -> bhlm', x, y)
+        return ret
+
+    def _get_relative_embeddings(self, relative_embeddings, length):
+        max_relative_position = 2 * self.window_size + 1
+        # Pad first before slice to avoid using cond ops.
+        pad_length = max(length - (self.window_size + 1), 0)
+        slice_start_position = max((self.window_size + 1) - length, 0)
+        slice_end_position = slice_start_position + 2 * length - 1
+        if pad_length > 0:
+            padded_relative_embeddings = F.pad(
+                relative_embeddings,
+                commons.convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]))
+        else:
+            padded_relative_embeddings = relative_embeddings
+        used_relative_embeddings = padded_relative_embeddings[
+            :, slice_start_position:slice_end_position
+        ]
+        return used_relative_embeddings
+
+    def _relative_position_to_absolute_position(self, x):
+        """
+        x: [b, h, l, 2*l-1]
+        ret: [b, h, l, l]
+        """
+        batch, heads, length, _ = x.size()
+        # Concat columns of pad to shift from relative to absolute indexing.
+        x = F.pad(x, commons.convert_pad_shape(
+            [[0, 0], [0, 0], [0, 0], [0, 1]]
+        ))
+
+        # Concat extra elements so to add up to shape (len+1, 2*len-1).
+        x_flat = x.view([batch, heads, length * 2 * length])
+        x_flat = F.pad(x_flat, commons.convert_pad_shape(
+            [[0, 0], [0, 0], [0, length-1]]
+        ))
+
+        # Reshape and slice out the padded elements.
+        x_final = x_flat.view([batch, heads, length+1, 2*length-1])[
+            :, :, :length, length-1:
+        ]
+        return x_final
+
+    def _absolute_position_to_relative_position(self, x):
+        """
+        x: [b, h, l, l]
+        ret: [b, h, l, 2*l-1]
+        """
+        batch, heads, length, _ = x.size()
+        # padd along column
+        x = F.pad(x, commons.convert_pad_shape(
+            [[0, 0], [0, 0], [0, 0], [0, length-1]]
+        ))
+        x_flat = x.view([batch, heads, length**2 + length*(length - 1)])
+        # add 0's in the beginning that will skew the elements after reshape
+        x_flat = F.pad(x_flat, commons.convert_pad_shape(
+            [[0, 0], [0, 0], [length, 0]]
+        ))
+        x_final = x_flat.view([batch, heads, length, 2*length])[:, :, :, 1:]
+        return x_final
+
+    def _attention_bias_proximal(self, length):
+        """Bias for self-attention to encourage attention to close positions.
+        Args:
+          length: an integer scalar.
+        Returns:
+          a Tensor with shape [1, 1, length, length]
+        """
+        r = torch.arange(length, dtype=torch.float32)
+        diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
+        return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
+
+
+class FFN(nn.Module):
+    def __init__(
+        self, 
+        in_channels, 
+        out_channels, 
+        filter_channels, 
+        kernel_size, 
+        p_dropout=0., 
+        activation=None, 
+        causal=False
+    ):
+        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.activation = activation
+        self.causal = causal
+
+        if causal:
+            self.padding = self._causal_padding
+        else:
+            self.padding = self._same_padding
+
+        self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
+        self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
+        self.drop = nn.Dropout(p_dropout)
+
+    def forward(self, x, x_mask):
+        x = self.conv_1(self.padding(x * x_mask))
+        if self.activation == "gelu":
+            x = x * torch.sigmoid(1.702 * x)
+        else:
+            x = torch.relu(x)
+        x = self.drop(x)
+        x = self.conv_2(self.padding(x * x_mask))
+        return x * x_mask
+
+    def _causal_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l = self.kernel_size - 1
+        pad_r = 0
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, commons.convert_pad_shape(padding))
+        return x
+
+    def _same_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l = (self.kernel_size - 1) // 2
+        pad_r = self.kernel_size // 2
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, commons.convert_pad_shape(padding))
+        return x

commons.py

@@ -0,0 +1,181 @@
+# from https://github.com/jaywalnut310/vits
+import math
+import torch
+from torch.nn import functional as F
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+def convert_pad_shape(pad_shape):
+    l = pad_shape[::-1]
+    pad_shape = [item for sublist in l for item in sublist]
+    return pad_shape
+
+
+def intersperse(lst, item):
+    result = [item] * (len(lst) * 2 + 1)
+    result[1::2] = lst
+    return result
+
+
+def kl_divergence(m_p, logs_p, m_q, logs_q):
+    """KL(P||Q)"""
+    kl = (logs_q - logs_p) - 0.5
+    kl += 0.5 * (torch.exp(2. * logs_p) + ((m_p - m_q)**2)) * torch.exp(-2. * logs_q)
+    return kl
+
+
+def rand_gumbel(shape):
+    """Sample from the Gumbel distribution, protect from overflows."""
+    uniform_samples = torch.rand(shape) * 0.99998 + 0.00001
+    return -torch.log(-torch.log(uniform_samples))
+
+
+def rand_gumbel_like(x):
+    g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device)
+    return g
+
+
+def slice_segments(x, ids_str, segment_size=4):
+    ret = torch.zeros_like(x[:, :, :segment_size])
+    for i in range(x.size(0)):
+        idx_str = ids_str[i]
+        idx_end = idx_str + segment_size
+        ret[i] = x[i, :, idx_str:idx_end]
+    return ret
+
+
+def rand_slice_segments(x, x_lengths=None, segment_size=4):
+    b, d, t = x.size()
+    if x_lengths is None:
+        x_lengths = t
+    ids_str_max = x_lengths - segment_size + 1
+    ids_str = (torch.rand([b]).to(device=x.device)
+               * ids_str_max).to(dtype=torch.long)
+    ids_str = torch.max(torch.zeros(ids_str.size()).to(ids_str.device), ids_str).to(dtype=torch.long)
+    ret = slice_segments(x, ids_str, segment_size)
+    return ret, ids_str
+
+def rand_slice_segments_for_cat(x, x_lengths=None, segment_size=4):
+    b, d, t = x.size()
+    if x_lengths is None:
+        x_lengths = t
+    ids_str_max = x_lengths - segment_size + 1
+    ids_str = torch.rand([b//2]).to(device=x.device)
+    ids_str = (torch.cat([ids_str,ids_str], dim=0)
+               * ids_str_max).to(dtype=torch.long)
+    ids_str = torch.max(torch.zeros(ids_str.size()).to(ids_str.device), ids_str).to(dtype=torch.long)
+    ret = slice_segments(x, ids_str, segment_size)
+    return ret, ids_str
+
+
+
+
+def get_timing_signal_1d(
+        length, channels, min_timescale=1.0, max_timescale=1.0e4):
+    position = torch.arange(length, dtype=torch.float)
+    num_timescales = channels // 2
+    log_timescale_increment = (
+        math.log(float(max_timescale) / float(min_timescale)) / (num_timescales - 1)
+    )
+    inv_timescales = min_timescale * torch.exp(
+        torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment
+    )
+    scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1)
+    signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0)
+    signal = F.pad(signal, [0, 0, 0, channels % 2])
+    signal = signal.view(1, channels, length)
+    return signal
+
+
+def add_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4):
+    b, channels, length = x.size()
+    signal = get_timing_signal_1d(
+        length, channels, min_timescale, max_timescale
+    )
+    return x + signal.to(dtype=x.dtype, device=x.device)
+
+
+def cat_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4, axis=1):
+    b, channels, length = x.size()
+    signal = get_timing_signal_1d(
+        length, channels, min_timescale, max_timescale
+    )
+    return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis)
+
+
+def subsequent_mask(length):
+    mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
+    return mask
+
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
+
+
+def convert_pad_shape(pad_shape):
+    l = pad_shape[::-1]
+    pad_shape = [item for sublist in l for item in sublist]
+    return pad_shape
+
+
+def shift_1d(x):
+    x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]
+    return x
+
+
+def sequence_mask(length, max_length=None):
+    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 generate_path(duration, mask):
+    """
+    duration: [b, 1, t_x]
+    mask: [b, 1, t_y, t_x]
+    """
+    device = duration.device
+
+    b, _, t_y, t_x = mask.shape
+    cum_duration = torch.cumsum(duration, -1)
+
+    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 - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
+    path = path.unsqueeze(1).transpose(2, 3) * mask
+    return path
+
+
+def clip_grad_value_(parameters, clip_value, norm_type=2):
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+    norm_type = float(norm_type)
+    if clip_value is not None:
+        clip_value = float(clip_value)
+
+    total_norm = 0
+    for p in parameters:
+        param_norm = p.grad.data.norm(norm_type)
+        total_norm += param_norm.item() ** norm_type
+        if clip_value is not None:
+            p.grad.data.clamp_(min=-clip_value, max=clip_value)
+    total_norm = total_norm ** (1. / norm_type)
+    return total_norm

configs/config_en.yaml

@@ -0,0 +1,68 @@
+train:
+  log_interval: 200   # step unit
+  eval_interval: 400  # step unit
+  save_interval: 50  # epoch unit: 50 for baseline / 500 for fine-tuning
+  seed: 1234
+  epochs: 7000
+  learning_rate: 2e-4 
+  betas: [0.8, 0.99]
+  eps: 1e-9
+  batch_size: 48
+  fp16_run: True  #False
+  lr_decay: 0.999875
+  segment_size: 8192
+  c_mel: 45
+  c_kl: 1.0
+  c_vq: 1.
+  c_commit: 0.2
+  c_yin: 45.
+  log_path: "/pits/logs"
+  n_sample: 3
+  alpha: 200
+
+data:
+  data_path: "/DATA/audio/VCTK-0.92"
+  training_files: "filelists/vctk_train_g2p.txt"
+  validation_files: "filelists/vctk_val_g2p.txt"
+  languages: "en_US" 
+  text_cleaners: ["english_cleaners"]
+  sampling_rate: 22050
+  filter_length: 1024
+  hop_length: 256
+  win_length: 1024
+  n_mel_channels: 80
+  mel_fmin: 0.0
+  mel_fmax: null
+  add_blank: True
+  speakers: ["p225", "p226", "p227", "p228", "p229", "p230", "p231", "p232", "p233", "p234", "p236", "p237", "p238", "p239", "p240", "p241", "p243", "p244", "p245", "p246", "p247", "p248", "p249", "p250", "p251", "p252", "p253", "p254", "p255", "p256", "p257", "p258", "p259", "p260", "p261", "p262", "p263", "p264", "p265", "p266", "p267", "p268", "p269", "p270", "p271", "p272", "p273", "p274", "p275", "p276", "p277", "p278", "p279", "p281", "p282", "p283", "p284", "p285", "p286", "p287", "p288", "p292", "p293", "p294", "p295", "p297", "p298", "p299", "p300", "p301", "p302", "p303", "p304", "p305", "p306", "p307", "p308", "p310", "p311", "p312", "p313", "p314", "p316", "p317", "p318", "p323", "p326", "p329", "p330", "p333", "p334", "p335", "p336", "p339", "p340", "p341", "p343", "p345", "p347", "p351", "p360", "p361", "p362", "p363", "p364", "p374", "p376", "s5"]
+  persistent_workers: True
+  midi_start: -5
+  midi_end: 75
+  midis: 80
+  ying_window: 2048
+  ying_hop: 256
+  tau_max: 2048
+  octave_range: 24
+  
+model:
+  inter_channels: 192
+  hidden_channels: 192
+  filter_channels: 768
+  n_heads: 2
+  n_layers: 6
+  kernel_size: 3
+  p_dropout: 0.1
+  resblock: "1"
+  resblock_kernel_sizes: [3,7,11]
+  resblock_dilation_sizes: [[1,3,5], [1,3,5], [1,3,5]]
+  upsample_rates: [8,8,2,2]
+  upsample_initial_channel: 512
+  upsample_kernel_sizes: [16,16,4,4]
+  n_layers_q: 3
+  use_spectral_norm: False
+  gin_channels: 256
+  codebook_size: 320
+  yin_channels: 80
+  yin_start: 15 # scope start bin in nansy = 1.5/8 
+  yin_scope: 50 # scope ratio in nansy = 5/8 
+  yin_shift_range: 15 # same as default start index of yingram

logs/pits_vctk_AD_3000.pth

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

models.py

@@ -0,0 +1,1383 @@
+# from https://github.com/jaywalnut310/vits
+# from https://github.com/ncsoft/avocodo
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn import Conv1d, ConvTranspose1d, Conv2d
+from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
+
+import modules
+import attentions
+import commons
+from commons import init_weights, get_padding
+#for Q option
+#from functions import vq, vq_st
+
+from analysis import Pitch
+from pqmf import PQMF
+
+
+class StochasticDurationPredictor(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 filter_channels,
+                 kernel_size,
+                 p_dropout,
+                 n_flows=4,
+                 gin_channels=0):
+        super().__init__()
+        # it needs to be removed from future version.
+        filter_channels = in_channels
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.log_flow = modules.Log()
+        self.flows = nn.ModuleList()
+        self.flows.append(modules.ElementwiseAffine(2))
+        for i in range(n_flows):
+            self.flows.append(
+                modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3))
+            self.flows.append(modules.Flip())
+
+        self.post_pre = nn.Conv1d(1, filter_channels, 1)
+        self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1)
+        self.post_convs = modules.DDSConv(filter_channels,
+                                          kernel_size,
+                                          n_layers=3,
+                                          p_dropout=p_dropout)
+        self.post_flows = nn.ModuleList()
+        self.post_flows.append(modules.ElementwiseAffine(2))
+        for i in range(4):
+            self.post_flows.append(
+                modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3))
+            self.post_flows.append(modules.Flip())
+
+        self.pre = nn.Conv1d(in_channels, filter_channels, 1)
+        self.proj = nn.Conv1d(filter_channels, filter_channels, 1)
+        self.convs = modules.DDSConv(filter_channels,
+                                     kernel_size,
+                                     n_layers=3,
+                                     p_dropout=p_dropout)
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
+
+    def forward(self,
+                x,
+                x_mask,
+                w=None,
+                g=None,
+                reverse=False,
+                noise_scale=1.0):
+        x = torch.detach(x)
+        x = self.pre(x)
+        if g is not None:
+            g = torch.detach(g)
+            x = x + self.cond(g)
+        x = self.convs(x, x_mask)
+        x = self.proj(x) * x_mask
+
+        if not reverse:
+            flows = self.flows
+            assert w is not None
+
+            logdet_tot_q = 0
+            h_w = self.post_pre(w)
+            h_w = self.post_convs(h_w, x_mask)
+            h_w = self.post_proj(h_w) * x_mask
+            e_q = torch.randn(w.size(0), 2, w.size(2)).to(
+                device=x.device, dtype=x.dtype) * x_mask
+            z_q = e_q
+            for flow in self.post_flows:
+                z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w))
+                logdet_tot_q += logdet_q
+            z_u, z1 = torch.split(z_q, [1, 1], 1)
+            u = torch.sigmoid(z_u) * x_mask
+            z0 = (w - u) * x_mask
+            logdet_tot_q += torch.sum(
+                (F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1, 2])
+            logq = torch.sum(
+                -0.5 * (math.log(2 * math.pi) +
+                        (e_q**2)) * x_mask, [1, 2]) - logdet_tot_q
+
+            logdet_tot = 0
+            z0, logdet = self.log_flow(z0, x_mask)
+            logdet_tot += logdet
+            z = torch.cat([z0, z1], 1)
+            for flow in flows:
+                z, logdet = flow(z, x_mask, g=x, reverse=reverse)
+                logdet_tot = logdet_tot + logdet
+            nll = torch.sum(0.5 * (math.log(2 * math.pi) +
+                                   (z**2)) * x_mask, [1, 2]) - logdet_tot
+            return nll + logq  # [b]
+        else:
+            flows = list(reversed(self.flows))
+            flows = flows[:-2] + [flows[-1]]  # remove a useless vflow
+            z = torch.randn(x.size(0), 2, x.size(2)).to(
+                device=x.device, dtype=x.dtype) * noise_scale
+            for flow in flows:
+                z = flow(z, x_mask, g=x, reverse=reverse)
+            z0, z1 = torch.split(z, [1, 1], 1)
+            logw = z0
+            return logw
+
+
+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 = modules.LayerNorm(filter_channels)
+        self.conv_2 = nn.Conv1d(filter_channels,
+                                filter_channels,
+                                kernel_size,
+                                padding=kernel_size // 2)
+        self.norm_2 = modules.LayerNorm(filter_channels)
+        self.proj = nn.Conv1d(filter_channels, 1, 1)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, in_channels, 1)
+
+    def forward(self, x, x_mask, g=None):
+        x = torch.detach(x)
+        if g is not None:
+            g = torch.detach(g)
+            x = x + self.cond(g)
+        x = self.conv_1(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_1(x)
+        x = self.drop(x)
+        x = self.conv_2(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_2(x)
+        x = self.drop(x)
+        x = self.proj(x * x_mask)
+        return x * x_mask
+
+
+class TextEncoder(nn.Module):
+
+    def __init__(self, n_vocab, out_channels, hidden_channels, filter_channels,
+                 n_heads, n_layers, kernel_size, p_dropout):
+        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.emb = nn.Embedding(n_vocab, hidden_channels)
+        nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)
+        self.emb_t = nn.Embedding(6, hidden_channels)
+        nn.init.normal_(self.emb_t.weight, 0.0, hidden_channels**-0.5)
+
+        self.encoder = attentions.Encoder(hidden_channels, filter_channels,
+                                          n_heads, n_layers, kernel_size,
+                                          p_dropout)
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(self, x, t, x_lengths):
+        t_zero = (t == 0)
+        emb_t = self.emb_t(t)
+        emb_t[t_zero, :] = 0
+        x = (self.emb(x) + emb_t) * math.sqrt(
+            self.hidden_channels)  # [b, t, h]
+        #x = torch.transpose(x, 1, -1)  # [b, h, t]
+        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(1)),
+                                 1).to(x.dtype)
+        #x = self.encoder(x * x_mask, x_mask)
+        x = torch.einsum('btd,but->bdt', x, x_mask)
+        x = self.encoder(x, x_mask)
+        stats = self.proj(x) * x_mask
+
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        return x, m, logs, x_mask
+
+
+class ResidualCouplingBlock(nn.Module):
+
+    def __init__(self,
+                 channels,
+                 hidden_channels,
+                 kernel_size,
+                 dilation_rate,
+                 n_layers,
+                 n_flows=4,
+                 gin_channels=0):
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.flows = nn.ModuleList()
+        for i in range(n_flows):
+            self.flows.append(
+                modules.ResidualCouplingLayer(channels,
+                                              hidden_channels,
+                                              kernel_size,
+                                              dilation_rate,
+                                              n_layers,
+                                              gin_channels=gin_channels,
+                                              mean_only=True))
+            self.flows.append(modules.Flip())
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        if not reverse:
+            for flow in self.flows:
+                x, _ = flow(x, x_mask, g=g, reverse=reverse)
+        else:
+            for flow in reversed(self.flows):
+                x = flow(x, x_mask, g=g, reverse=reverse)
+        return x
+
+
+class PosteriorEncoder(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 hidden_channels,
+                 kernel_size,
+                 dilation_rate,
+                 n_layers,
+                 gin_channels=0):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+
+        self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
+        self.enc = modules.WN(hidden_channels,
+                              kernel_size,
+                              dilation_rate,
+                              n_layers,
+                              gin_channels=gin_channels)
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(self, x, x_lengths, g=None):
+        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)),
+                                 1).to(x.dtype)
+        x = self.pre(x) * x_mask
+        x = self.enc(x, x_mask, g=g)
+        stats = self.proj(x) * x_mask
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask
+        return z, m, logs, x_mask
+
+
+class Generator(nn.Module):
+
+    def __init__(self,
+                 initial_channel,
+                 resblock,
+                 resblock_kernel_sizes,
+                 resblock_dilation_sizes,
+                 upsample_rates,
+                 upsample_initial_channel,
+                 upsample_kernel_sizes,
+                 gin_channels=0):
+        super(Generator, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.conv_pre = Conv1d(initial_channel,
+                               upsample_initial_channel,
+                               7,
+                               1,
+                               padding=3)
+        resblock = modules.ResBlock1 if resblock == '1' else modules.ResBlock2
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(
+                weight_norm(
+                    ConvTranspose1d(upsample_initial_channel // (2**i),
+                                    upsample_initial_channel // (2**(i + 1)),
+                                    k,
+                                    u,
+                                    padding=(k - u) // 2)))
+
+        self.resblocks = nn.ModuleList()
+        self.conv_posts = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel // (2**(i + 1))
+            for j, (k, d) in enumerate(
+                    zip(resblock_kernel_sizes, resblock_dilation_sizes)):
+                self.resblocks.append(resblock(ch, k, d))
+            if i >= len(self.ups) - 3:
+                self.conv_posts.append(
+                    Conv1d(ch, 1, 7, 1, padding=3, bias=False))
+        self.ups.apply(init_weights)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+
+    def forward(self, x, g=None):
+        x = self.conv_pre(x)
+        if g is not None:
+            x = x + self.cond(g)
+
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                xs = xs + self.resblocks[i * self.num_kernels + j](x) if xs is not None \
+                     else self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_posts[-1](x)
+        x = torch.tanh(x)
+
+        return x
+
+    def hier_forward(self, x, g=None):
+        outs = []
+        x = self.conv_pre(x)
+        if g is not None:
+            x = x + self.cond(g)
+
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                xs = xs + self.resblocks[i * self.num_kernels + j](x) if xs is not None \
+                     else self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+            if i >= self.num_upsamples - 3:
+                _x = F.leaky_relu(x)
+                _x = self.conv_posts[i - self.num_upsamples + 3](_x)
+                _x = torch.tanh(_x)
+                outs.append(_x)
+        return outs
+
+    def remove_weight_norm(self):
+        print('Removing weight norm...')
+        for l in self.ups:
+            remove_weight_norm(l)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+
+
+class DiscriminatorP(nn.Module):
+
+    def __init__(self,
+                 period,
+                 kernel_size=5,
+                 stride=3,
+                 use_spectral_norm=False):
+        super(DiscriminatorP, self).__init__()
+        self.period = period
+        self.use_spectral_norm = use_spectral_norm
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.convs = nn.ModuleList([
+            norm_f(
+                Conv2d(1,
+                       32, (kernel_size, 1), (stride, 1),
+                       padding=(get_padding(kernel_size, 1), 0))),
+            norm_f(
+                Conv2d(32,
+                       128, (kernel_size, 1), (stride, 1),
+                       padding=(get_padding(kernel_size, 1), 0))),
+            norm_f(
+                Conv2d(128,
+                       512, (kernel_size, 1), (stride, 1),
+                       padding=(get_padding(kernel_size, 1), 0))),
+            norm_f(
+                Conv2d(512,
+                       1024, (kernel_size, 1), (stride, 1),
+                       padding=(get_padding(kernel_size, 1), 0))),
+            norm_f(
+                Conv2d(1024,
+                       1024, (kernel_size, 1),
+                       1,
+                       padding=(get_padding(kernel_size, 1), 0))),
+        ])
+        self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
+
+    def forward(self, x):
+        fmap = []
+
+        # 1d to 2d
+        b, c, t = x.shape
+        if t % self.period != 0:  # pad first
+            n_pad = self.period - (t % self.period)
+            x = F.pad(x, (0, n_pad), "reflect")
+            t = t + n_pad
+        x = x.view(b, c, t // self.period, self.period)
+
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class DiscriminatorS(nn.Module):
+
+    def __init__(self, use_spectral_norm=False):
+        super(DiscriminatorS, self).__init__()
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.convs = nn.ModuleList([
+            norm_f(Conv1d(1, 16, 15, 1, padding=7)),
+            norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)),
+            norm_f(Conv1d(64, 256, 41, 4, groups=16, padding=20)),
+            norm_f(Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
+            norm_f(Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
+            norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
+        ])
+        self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
+
+    def forward(self, x):
+        fmap = []
+
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class MultiPeriodDiscriminator(nn.Module):
+
+    def __init__(self, use_spectral_norm=False):
+        super(MultiPeriodDiscriminator, self).__init__()
+        periods = [2, 3, 5, 7, 11]
+
+        discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)]
+        discs = discs + \
+            [DiscriminatorP(i, use_spectral_norm=use_spectral_norm)
+             for i in periods]
+        self.discriminators = nn.ModuleList(discs)
+
+    def forward(self, y, y_hat):
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            y_d_gs.append(y_d_g)
+            fmap_rs.append(fmap_r)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+##### Avocodo
+class CoMBDBlock(torch.nn.Module):
+
+    def __init__(
+            self,
+            h_u,  # List[int],
+            d_k,  # List[int],
+            d_s,  # List[int],
+            d_d,  # List[int],
+            d_g,  # List[int],
+            d_p,  # List[int],
+            op_f,  # int,
+            op_k,  # int,
+            op_g,  # int,
+            use_spectral_norm=False):
+        super(CoMBDBlock, self).__init__()
+        norm_f = weight_norm if use_spectral_norm is False else spectral_norm
+
+        self.convs = nn.ModuleList()
+        filters = [[1, h_u[0]]]
+        for i in range(len(h_u) - 1):
+            filters.append([h_u[i], h_u[i + 1]])
+        for _f, _k, _s, _d, _g, _p in zip(filters, d_k, d_s, d_d, d_g, d_p):
+            self.convs.append(
+                norm_f(
+                    Conv1d(in_channels=_f[0],
+                           out_channels=_f[1],
+                           kernel_size=_k,
+                           stride=_s,
+                           dilation=_d,
+                           groups=_g,
+                           padding=_p)))
+        self.projection_conv = norm_f(
+            Conv1d(in_channels=filters[-1][1],
+                   out_channels=op_f,
+                   kernel_size=op_k,
+                   groups=op_g))
+
+    def forward(self, x, b_y, b_y_hat):
+        fmap_r = []
+        fmap_g = []
+        for block in self.convs:
+            x = block(x)
+            x = F.leaky_relu(x, 0.2)
+            f_r, f_g = x.split([b_y, b_y_hat], dim=0)
+            fmap_r.append(f_r.tile([2, 1, 1]) if b_y < b_y_hat else f_r)
+            fmap_g.append(f_g)
+        x = self.projection_conv(x)
+        x_r, x_g = x.split([b_y, b_y_hat], dim=0)
+        return x_r.tile([2, 1, 1
+                         ]) if b_y < b_y_hat else x_r, x_g, fmap_r, fmap_g
+
+
+class CoMBD(torch.nn.Module):
+
+    def __init__(self, use_spectral_norm=False):
+        super(CoMBD, self).__init__()
+        self.pqmf_list = nn.ModuleList([
+            PQMF(4, 192, 0.13, 10.0),  #lv2
+            PQMF(2, 256, 0.25, 10.0)  #lv1
+        ])
+        combd_h_u = [[16, 64, 256, 1024, 1024, 1024] for _ in range(3)]
+        combd_d_k = [[7, 11, 11, 11, 11, 5], [11, 21, 21, 21, 21, 5],
+                     [15, 41, 41, 41, 41, 5]]
+        combd_d_s = [[1, 1, 4, 4, 4, 1] for _ in range(3)]
+        combd_d_d = [[1, 1, 1, 1, 1, 1] for _ in range(3)]
+        combd_d_g = [[1, 4, 16, 64, 256, 1] for _ in range(3)]
+
+        combd_d_p = [[3, 5, 5, 5, 5, 2], [5, 10, 10, 10, 10, 2],
+                     [7, 20, 20, 20, 20, 2]]
+        combd_op_f = [1, 1, 1]
+        combd_op_k = [3, 3, 3]
+        combd_op_g = [1, 1, 1]
+
+        self.blocks = nn.ModuleList()
+        for _h_u, _d_k, _d_s, _d_d, _d_g, _d_p, _op_f, _op_k, _op_g in zip(
+                combd_h_u,
+                combd_d_k,
+                combd_d_s,
+                combd_d_d,
+                combd_d_g,
+                combd_d_p,
+                combd_op_f,
+                combd_op_k,
+                combd_op_g,
+        ):
+            self.blocks.append(
+                CoMBDBlock(
+                    _h_u,
+                    _d_k,
+                    _d_s,
+                    _d_d,
+                    _d_g,
+                    _d_p,
+                    _op_f,
+                    _op_k,
+                    _op_g,
+                ))
+
+    def _block_forward(self, ys, ys_hat, blocks):
+        outs_real = []
+        outs_fake = []
+        f_maps_real = []
+        f_maps_fake = []
+        for y, y_hat, block in zip(ys, ys_hat,
+                                   blocks):  #y:B, y_hat: 2B if i!=-1 else B,B
+            b_y = y.shape[0]
+            b_y_hat = y_hat.shape[0]
+            cat_y = torch.cat([y, y_hat], dim=0)
+            out_real, out_fake, f_map_r, f_map_g = block(cat_y, b_y, b_y_hat)
+            outs_real.append(out_real)
+            outs_fake.append(out_fake)
+            f_maps_real.append(f_map_r)
+            f_maps_fake.append(f_map_g)
+        return outs_real, outs_fake, f_maps_real, f_maps_fake
+
+    def _pqmf_forward(self, ys, ys_hat):
+        # preprocess for multi_scale forward
+        multi_scale_inputs_hat = []
+        for pqmf_ in self.pqmf_list:
+            multi_scale_inputs_hat.append(pqmf_.analysis(ys_hat[-1])[:, :1, :])
+
+        # real
+        # for hierarchical forward
+        #outs_real_, f_maps_real_ = self._block_forward(
+        #    ys, self.blocks)
+
+        # for multi_scale forward
+        #outs_real, f_maps_real = self._block_forward(
+        #        ys[:-1], self.blocks[:-1], outs_real, f_maps_real)
+        #outs_real.extend(outs_real[:-1])
+        #f_maps_real.extend(f_maps_real[:-1])
+
+        #outs_real = [torch.cat([o,o], dim=0) if i!=len(outs_real_)-1 else o for i,o in enumerate(outs_real_)]
+        #f_maps_real = [[torch.cat([fmap,fmap], dim=0) if i!=len(f_maps_real_)-1 else fmap for fmap in fmaps ] \
+        #        for i,fmaps in enumerate(f_maps_real_)]
+
+        inputs_fake = [
+            torch.cat([y, multi_scale_inputs_hat[i]], dim=0)
+            if i != len(ys_hat) - 1 else y for i, y in enumerate(ys_hat)
+        ]
+        outs_real, outs_fake, f_maps_real, f_maps_fake = self._block_forward(
+            ys, inputs_fake, self.blocks)
+
+        # predicted
+        # for hierarchical forward
+        #outs_fake, f_maps_fake = self._block_forward(
+        #    inputs_fake, self.blocks)
+
+        #outs_real_, f_maps_real_ = self._block_forward(
+        #    ys, self.blocks)
+        # for multi_scale forward
+        #outs_fake, f_maps_fake = self._block_forward(
+        #    multi_scale_inputs_hat, self.blocks[:-1], outs_fake, f_maps_fake)
+
+        return outs_real, outs_fake, f_maps_real, f_maps_fake
+
+    def forward(self, ys, ys_hat):
+        outs_real, outs_fake, f_maps_real, f_maps_fake = self._pqmf_forward(
+            ys, ys_hat)
+        return outs_real, outs_fake, f_maps_real, f_maps_fake
+
+
+class MDC(torch.nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 strides,
+                 kernel_size,
+                 dilations,
+                 use_spectral_norm=False):
+        super(MDC, self).__init__()
+        norm_f = weight_norm if not use_spectral_norm else spectral_norm
+        self.d_convs = nn.ModuleList()
+        for _k, _d in zip(kernel_size, dilations):
+            self.d_convs.append(
+                norm_f(
+                    Conv1d(in_channels=in_channels,
+                           out_channels=out_channels,
+                           kernel_size=_k,
+                           dilation=_d,
+                           padding=get_padding(_k, _d))))
+        self.post_conv = norm_f(
+            Conv1d(in_channels=out_channels,
+                   out_channels=out_channels,
+                   kernel_size=3,
+                   stride=strides,
+                   padding=get_padding(_k, _d)))
+        self.softmax = torch.nn.Softmax(dim=-1)
+
+    def forward(self, x):
+        _out = None
+        for _l in self.d_convs:
+            _x = torch.unsqueeze(_l(x), -1)
+            _x = F.leaky_relu(_x, 0.2)
+            _out = torch.cat([_out, _x], axis=-1) if _out is not None \
+                   else _x
+        x = torch.sum(_out, dim=-1)
+        x = self.post_conv(x)
+        x = F.leaky_relu(x, 0.2)  # @@
+
+        return x
+
+
+class SBDBlock(torch.nn.Module):
+
+    def __init__(self,
+                 segment_dim,
+                 strides,
+                 filters,
+                 kernel_size,
+                 dilations,
+                 use_spectral_norm=False):
+        super(SBDBlock, self).__init__()
+        norm_f = weight_norm if not use_spectral_norm else spectral_norm
+        self.convs = nn.ModuleList()
+        filters_in_out = [(segment_dim, filters[0])]
+        for i in range(len(filters) - 1):
+            filters_in_out.append([filters[i], filters[i + 1]])
+
+        for _s, _f, _k, _d in zip(strides, filters_in_out, kernel_size,
+                                  dilations):
+            self.convs.append(
+                MDC(in_channels=_f[0],
+                    out_channels=_f[1],
+                    strides=_s,
+                    kernel_size=_k,
+                    dilations=_d,
+                    use_spectral_norm=use_spectral_norm))
+        self.post_conv = norm_f(
+            Conv1d(in_channels=_f[1],
+                   out_channels=1,
+                   kernel_size=3,
+                   stride=1,
+                   padding=3 // 2))  # @@
+
+    def forward(self, x):
+        fmap_r = []
+        fmap_g = []
+        for _l in self.convs:
+            x = _l(x)
+            f_r, f_g = torch.chunk(x, 2, dim=0)
+            fmap_r.append(f_r)
+            fmap_g.append(f_g)
+        x = self.post_conv(x)  # @@
+        x_r, x_g = torch.chunk(x, 2, dim=0)
+        return x_r, x_g, fmap_r, fmap_g
+
+
+class MDCDConfig:
+
+    def __init__(self):
+        self.pqmf_params = [16, 256, 0.03, 10.0]
+        self.f_pqmf_params = [64, 256, 0.1, 9.0]
+        self.filters = [[64, 128, 256, 256, 256], [64, 128, 256, 256, 256],
+                        [64, 128, 256, 256, 256], [32, 64, 128, 128, 128]]
+        self.kernel_sizes = [[[7, 7, 7], [7, 7, 7], [7, 7, 7], [7, 7, 7],
+                              [7, 7, 7]],
+                             [[5, 5, 5], [5, 5, 5], [5, 5, 5], [5, 5, 5],
+                              [5, 5, 5]],
+                             [[3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3],
+                              [3, 3, 3]],
+                             [[5, 5, 5], [5, 5, 5], [5, 5, 5], [5, 5, 5],
+                              [5, 5, 5]]]
+        self.dilations = [[[5, 7, 11], [5, 7, 11], [5, 7, 11], [5, 7, 11],
+                           [5, 7, 11]],
+                          [[3, 5, 7], [3, 5, 7], [3, 5, 7], [3, 5, 7],
+                           [3, 5, 7]],
+                          [[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3],
+                           [1, 2, 3]],
+                          [[1, 2, 3], [1, 2, 3], [1, 2, 3], [2, 3, 5],
+                           [2, 3, 5]]]
+        self.strides = [[1, 1, 3, 3, 1], [1, 1, 3, 3, 1], [1, 1, 3, 3, 1],
+                        [1, 1, 3, 3, 1]]
+        self.band_ranges = [[0, 6], [0, 11], [0, 16], [0, 64]]
+        self.transpose = [False, False, False, True]
+        self.segment_size = 8192
+
+
+class SBD(torch.nn.Module):
+
+    def __init__(self, use_spectral_norm=False):
+        super(SBD, self).__init__()
+        self.config = MDCDConfig()
+        self.pqmf = PQMF(*self.config.pqmf_params)
+        if True in self.config.transpose:
+            self.f_pqmf = PQMF(*self.config.f_pqmf_params)
+        else:
+            self.f_pqmf = None
+
+        self.discriminators = torch.nn.ModuleList()
+
+        for _f, _k, _d, _s, _br, _tr in zip(self.config.filters,
+                                            self.config.kernel_sizes,
+                                            self.config.dilations,
+                                            self.config.strides,
+                                            self.config.band_ranges,
+                                            self.config.transpose):
+            if _tr:
+                segment_dim = self.config.segment_size // _br[1] - _br[0]
+            else:
+                segment_dim = _br[1] - _br[0]
+
+            self.discriminators.append(
+                SBDBlock(segment_dim=segment_dim,
+                         filters=_f,
+                         kernel_size=_k,
+                         dilations=_d,
+                         strides=_s,
+                         use_spectral_norm=use_spectral_norm))
+
+    def forward(self, y, y_hat):
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        y_in = self.pqmf.analysis(y)
+        y_hat_in = self.pqmf.analysis(y_hat)
+        y_in_f = self.f_pqmf.analysis(y)
+        y_hat_in_f = self.f_pqmf.analysis(y_hat)
+
+        for d, br, tr in zip(self.discriminators, self.config.band_ranges,
+                             self.config.transpose):
+            if not tr:
+                _y_in = y_in[:, br[0]:br[1], :]
+                _y_hat_in = y_hat_in[:, br[0]:br[1], :]
+            else:
+                _y_in = y_in_f[:, br[0]:br[1], :]
+                _y_hat_in = y_hat_in_f[:, br[0]:br[1], :]
+                _y_in = torch.transpose(_y_in, 1, 2)
+                _y_hat_in = torch.transpose(_y_hat_in, 1, 2)
+            #y_d_r, fmap_r = d(_y_in)
+            #y_d_g, fmap_g = d(_y_hat_in)
+            cat_y = torch.cat([_y_in, _y_hat_in], dim=0)
+            y_d_r, y_d_g, fmap_r, fmap_g = d(cat_y)
+            y_d_rs.append(y_d_r)
+            fmap_rs.append(fmap_r)
+            y_d_gs.append(y_d_g)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class AvocodoDiscriminator(nn.Module):
+
+    def __init__(self, use_spectral_norm=False):
+        super(AvocodoDiscriminator, self).__init__()
+        self.combd = CoMBD(use_spectral_norm)
+        self.sbd = SBD(use_spectral_norm)
+
+    def forward(self, y, ys_hat):
+        ys = [
+            self.combd.pqmf_list[0].analysis(y)[:, :1],  #lv2
+            self.combd.pqmf_list[1].analysis(y)[:, :1],  #lv1
+            y
+        ]
+        y_c_rs, y_c_gs, fmap_c_rs, fmap_c_gs = self.combd(ys, ys_hat)
+        y_s_rs, y_s_gs, fmap_s_rs, fmap_s_gs = self.sbd(y, ys_hat[-1])
+        y_c_rs.extend(y_s_rs)
+        y_c_gs.extend(y_s_gs)
+        fmap_c_rs.extend(fmap_s_rs)
+        fmap_c_gs.extend(fmap_s_gs)
+        return y_c_rs, y_c_gs, fmap_c_rs, fmap_c_gs
+
+
+##### Avocodo
+
+
+class YingDecoder(nn.Module):
+
+    def __init__(self,
+                 hidden_channels,
+                 kernel_size,
+                 dilation_rate,
+                 n_layers,
+                 yin_start,
+                 yin_scope,
+                 yin_shift_range,
+                 gin_channels=0):
+        super().__init__()
+        self.in_channels = yin_scope
+        self.out_channels = yin_scope
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+
+        self.yin_start = yin_start
+        self.yin_scope = yin_scope
+        self.yin_shift_range = yin_shift_range
+
+        self.pre = nn.Conv1d(self.in_channels, hidden_channels, 1)
+        self.dec = modules.WN(hidden_channels,
+                              kernel_size,
+                              dilation_rate,
+                              n_layers,
+                              gin_channels=gin_channels)
+        self.proj = nn.Conv1d(hidden_channels, self.out_channels, 1)
+
+    def crop_scope(self, x, yin_start,
+                   scope_shift):  # x: tensor [B,C,T] #scope_shift: tensor [B]
+        return torch.stack([
+            x[i, yin_start + scope_shift[i]:yin_start + self.yin_scope +
+              scope_shift[i], :] for i in range(x.shape[0])
+        ],
+                           dim=0)
+
+    def infer(self, z_yin, z_mask, g=None):
+        B = z_yin.shape[0]
+        scope_shift = torch.randint(-self.yin_shift_range,
+                                    self.yin_shift_range, (B, ),
+                                    dtype=torch.int)
+        z_yin_crop = self.crop_scope(z_yin, self.yin_start, scope_shift)
+        x = self.pre(z_yin_crop) * z_mask
+        x = self.dec(x, z_mask, g=g)
+        yin_hat_crop = self.proj(x) * z_mask
+        return yin_hat_crop
+
+    def forward(self, z_yin, yin_gt, z_mask, g=None):
+        B = z_yin.shape[0]
+        scope_shift = torch.randint(-self.yin_shift_range,
+                                    self.yin_shift_range, (B, ),
+                                    dtype=torch.int)
+        z_yin_crop = self.crop_scope(z_yin, self.yin_start, scope_shift)
+        yin_gt_shifted_crop = self.crop_scope(yin_gt, self.yin_start,
+                                              scope_shift)
+        yin_gt_crop = self.crop_scope(yin_gt, self.yin_start,
+                                      torch.zeros_like(scope_shift))
+        x = self.pre(z_yin_crop) * z_mask
+        x = self.dec(x, z_mask, g=g)
+        yin_hat_crop = self.proj(x) * z_mask
+        return yin_gt_crop, yin_gt_shifted_crop, yin_hat_crop, z_yin_crop, scope_shift
+
+
+# For Q option
+#class VQEmbedding(nn.Module):
+#
+#    def __init__(self, codebook_size,
+#                 code_channels):
+#        super().__init__()
+#        self.embedding = nn.Embedding(codebook_size, code_channels)
+#        self.embedding.weight.data.uniform_(-1. / codebook_size,
+#                                            1. / codebook_size)
+#
+#    def forward(self, z_e_x):
+#        z_e_x_ = z_e_x.permute(0, 2, 1).contiguous()
+#        latent_indices = vq(z_e_x_, self.embedding.weight)
+#        z_q = self.embedding(latent_indices).permute(0, 2, 1)
+#        return z_q
+#
+#    def straight_through(self, z_e_x):
+#        z_e_x_ = z_e_x.permute(0, 2, 1).contiguous()
+#        z_q_x_st_, indices = vq_st(z_e_x_, self.embedding.weight.detach())
+#        z_q_x_st = z_q_x_st_.permute(0, 2, 1).contiguous()
+#
+#        z_q_x_flatten = torch.index_select(self.embedding.weight,
+#                                           dim=0,
+#                                           index=indices)
+#        z_q_x_ = z_q_x_flatten.view_as(z_e_x_)
+#        z_q_x = z_q_x_.permute(0, 2, 1).contiguous()
+#        return z_q_x_st, z_q_x
+
+
+class SynthesizerTrn(nn.Module):
+    """
+    Synthesizer for Training
+    """
+
+    def __init__(
+            self,
+            n_vocab,
+            spec_channels,
+            segment_size,
+            midi_start,
+            midi_end,
+            octave_range,
+            inter_channels,
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            yin_channels,
+            yin_start,
+            yin_scope,
+            yin_shift_range,
+            n_speakers=0,
+            gin_channels=0,
+            use_sdp=True,
+            #codebook_size=256, #for Q option
+            **kwargs):
+
+        super().__init__()
+        self.n_vocab = n_vocab
+        self.spec_channels = spec_channels
+        self.inter_channels = inter_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.resblock = resblock
+        self.resblock_kernel_sizes = resblock_kernel_sizes
+        self.resblock_dilation_sizes = resblock_dilation_sizes
+        self.upsample_rates = upsample_rates
+        self.upsample_initial_channel = upsample_initial_channel
+        self.upsample_kernel_sizes = upsample_kernel_sizes
+        self.segment_size = segment_size
+        self.n_speakers = n_speakers
+        self.gin_channels = gin_channels
+
+        self.yin_channels = yin_channels
+        self.yin_start = yin_start
+        self.yin_scope = yin_scope
+
+        self.use_sdp = use_sdp
+        self.enc_p = TextEncoder(n_vocab, inter_channels, hidden_channels,
+                                 filter_channels, n_heads, n_layers,
+                                 kernel_size, p_dropout)
+        self.dec = Generator(
+            inter_channels - yin_channels +
+            yin_scope,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            gin_channels=gin_channels)
+
+        self.enc_spec = PosteriorEncoder(spec_channels,
+                                         inter_channels - yin_channels,
+                                         inter_channels - yin_channels,
+                                         5,
+                                         1,
+                                         16,
+                                         gin_channels=gin_channels)
+
+        self.enc_pitch = PosteriorEncoder(yin_channels,
+                                          yin_channels,
+                                          yin_channels,
+                                          5,
+                                          1,
+                                          16,
+                                          gin_channels=gin_channels)
+
+        self.flow = ResidualCouplingBlock(inter_channels,
+                                          hidden_channels,
+                                          5,
+                                          1,
+                                          4,
+                                          gin_channels=gin_channels)
+
+        if use_sdp:
+            self.dp = StochasticDurationPredictor(hidden_channels,
+                                                  192,
+                                                  3,
+                                                  0.5,
+                                                  4,
+                                                  gin_channels=gin_channels)
+        else:
+            self.dp = DurationPredictor(hidden_channels,
+                                        256,
+                                        3,
+                                        0.5,
+                                        gin_channels=gin_channels)
+
+        self.yin_dec = YingDecoder(yin_scope,
+                                   5,
+                                   1,
+                                   4,
+                                   yin_start,
+                                   yin_scope,
+                                   yin_shift_range,
+                                   gin_channels=gin_channels)
+
+        #self.vq = VQEmbedding(codebook_size, inter_channels - yin_channels)#inter_channels // 2)
+        self.emb_g = nn.Embedding(self.n_speakers, gin_channels)
+
+        self.pitch = Pitch(midi_start=midi_start,
+                           midi_end=midi_end,
+                           octave_range=octave_range)
+
+    def crop_scope(
+            self,
+            x,
+            scope_shift=0):  # x: list #need to modify for non-scalar shift
+        return [
+            i[:, self.yin_start + scope_shift:self.yin_start + self.yin_scope +
+              scope_shift, :] for i in x
+        ]
+
+    def crop_scope_tensor(
+            self, x,
+            scope_shift):  # x: tensor [B,C,T] #scope_shift: tensor [B]
+        return torch.stack([
+            x[i, self.yin_start + scope_shift[i]:self.yin_start +
+              self.yin_scope + scope_shift[i], :] for i in range(x.shape[0])
+        ],
+                           dim=0)
+
+    def yin_dec_infer(self, z_yin, z_mask, sid=None):
+        if self.n_speakers > 0:
+            g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
+        else:
+            g = None
+        return self.yin_dec.infer(z_yin, z_mask, g)
+
+    def forward(self,
+                x,
+                t,
+                x_lengths,
+                y,
+                y_lengths,
+                ying,
+                ying_lengths,
+                sid=None,
+                scope_shift=0):
+        x, m_p, logs_p, x_mask = self.enc_p(x, t, x_lengths)
+        if self.n_speakers > 0:
+            g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
+        else:
+            g = None
+
+        z_spec, m_spec, logs_spec, spec_mask = self.enc_spec(y, y_lengths, g=g)
+
+        #for Q option
+        #z_spec_q_st, z_spec_q = self.vq.straight_through(z_spec)
+        #z_spec_q_st = z_spec_q_st * spec_mask
+        #z_spec_q = z_spec_q * spec_mask
+
+        z_yin, m_yin, logs_yin, yin_mask = self.enc_pitch(ying, y_lengths, g=g)
+        z_yin_crop, logs_yin_crop, m_yin_crop = self.crop_scope(
+            [z_yin, logs_yin, m_yin], scope_shift)
+
+        #yin dec loss
+        yin_gt_crop, yin_gt_shifted_crop, yin_dec_crop, z_yin_crop_shifted, scope_shift = self.yin_dec(
+            z_yin, ying, yin_mask, g)
+
+        z = torch.cat([z_spec, z_yin], dim=1)
+        logs_q = torch.cat([logs_spec, logs_yin], dim=1)
+        m_q = torch.cat([m_spec, m_yin], dim=1)
+        y_mask = spec_mask
+
+        z_p = self.flow(z, y_mask, g=g)
+
+        z_dec = torch.cat([z_spec, z_yin_crop], dim=1)
+
+        z_dec_shifted = torch.cat([z_spec.detach(), z_yin_crop_shifted], dim=1)
+        z_dec_ = torch.cat([z_dec, z_dec_shifted], dim=0)
+
+        with torch.no_grad():
+            # negative cross-entropy
+            s_p_sq_r = torch.exp(-2 * logs_p)  # [b, d, t]
+            # [b, 1, t_s]
+            neg_cent1 = torch.sum(-0.5 * math.log(2 * math.pi) - logs_p, [1],
+                                  keepdim=True)
+            # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s], z_p: [b,d,t]
+            #neg_cent2 = torch.matmul(-0.5 * (z_p**2).transpose(1, 2), s_p_sq_r)
+            neg_cent2 = torch.einsum('bdt, bds -> bts', -0.5 * (z_p**2),
+                                     s_p_sq_r)
+            # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s]
+            #neg_cent3 = torch.matmul(z_p.transpose(1, 2), (m_p * s_p_sq_r))
+            neg_cent3 = torch.einsum('bdt, bds -> bts', z_p, (m_p * s_p_sq_r))
+            neg_cent4 = torch.sum(-0.5 * (m_p**2) * s_p_sq_r, [1],
+                                  keepdim=True)  # [b, 1, t_s]
+            neg_cent = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4
+
+            attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(
+                y_mask, -1)
+            from monotonic_align import maximum_path
+            attn = maximum_path(neg_cent,
+                                attn_mask.squeeze(1)).unsqueeze(1).detach()
+
+        w = attn.sum(2)
+        if self.use_sdp:
+            l_length = self.dp(x, x_mask, w, g=g)
+            l_length = l_length / torch.sum(x_mask)
+        else:
+            logw_ = torch.log(w + 1e-6) * x_mask
+            logw = self.dp(x, x_mask, g=g)
+            l_length = torch.sum(
+                (logw - logw_)**2, [1, 2]) / torch.sum(x_mask)  # for averaging
+
+        # expand prior
+        m_p = torch.einsum('bctn, bdn -> bdt', attn, m_p)
+        logs_p = torch.einsum('bctn, bdn -> bdt', attn, logs_p)
+
+        #z_slice, ids_slice = commons.rand_slice_segments(z_dec, y_lengths, self.segment_size)
+        #o = self.dec(z_slice, g=g)
+        z_slice, ids_slice = commons.rand_slice_segments_for_cat(
+            z_dec_, torch.cat([y_lengths, y_lengths], dim=0),
+            self.segment_size)
+        o_ = self.dec.hier_forward(z_slice, g=torch.cat([g, g], dim=0))
+        o = [torch.chunk(o_hier, 2, dim=0)[0] for o_hier in o_]
+
+        o_pad = F.pad(o_[-1], (768, 768 + (-o_[-1].shape[-1]) % 256 + 256 *
+                               (o_[-1].shape[-1] % 256 == 0)),
+                      mode='constant').squeeze(1)
+        yin_hat = self.pitch.yingram(o_pad)
+        yin_hat_crop = self.crop_scope([yin_hat])[0]
+        yin_hat_shifted = self.crop_scope_tensor(
+            torch.chunk(yin_hat, 2, dim=0)[0], scope_shift)
+        return o, l_length, attn, ids_slice, x_mask, y_mask, o_, \
+               (z, z_p, m_p, logs_p, m_q, logs_q), \
+               (z_dec_), \
+               (z_spec, m_spec, logs_spec, spec_mask, z_yin, m_yin, logs_yin, yin_mask), \
+               (yin_gt_crop, yin_gt_shifted_crop, yin_dec_crop, yin_hat_crop, scope_shift, yin_hat_shifted)
+
+    def infer(self,
+              x,
+              t,
+              x_lengths,
+              sid=None,
+              noise_scale=1,
+              length_scale=1,
+              noise_scale_w=1.,
+              max_len=None,
+              scope_shift=0):  #need to fix #vector scope shift needed
+        x, m_p, logs_p, x_mask = self.enc_p(x, t, x_lengths)
+        if self.n_speakers > 0:
+            g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
+        else:
+            g = None
+
+        if self.use_sdp:
+            logw = self.dp(x,
+                           x_mask,
+                           g=g,
+                           reverse=True,
+                           noise_scale=noise_scale_w)
+        else:
+            logw = self.dp(x, x_mask, g=g)
+        w = torch.exp(logw) * x_mask * length_scale
+        w_ceil = torch.ceil(w)
+        y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
+        y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, None),
+                                 1).to(x_mask.dtype)
+        attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
+        attn = commons.generate_path(w_ceil, attn_mask)
+
+        m_p = torch.einsum('bctn, bdn -> bdt', attn, m_p)
+        logs_p = torch.einsum('bctn, bdn -> bdt', attn, logs_p)
+
+        z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
+        z = self.flow(z_p, y_mask, g=g, reverse=True)
+        z_spec, z_yin = torch.split(z,
+                                    self.inter_channels - self.yin_channels,
+                                    dim=1)
+        z_yin_crop = self.crop_scope([z_yin], scope_shift)[0]
+        z_crop = torch.cat([z_spec, z_yin_crop], dim=1)
+        o = self.dec((z_crop * y_mask)[:, :, :max_len], g=g)
+        return o, attn, y_mask, (z_crop, z, z_p, m_p, logs_p)
+
+    def infer_pre_decoder(self,
+                          x,
+                          t,
+                          x_lengths,
+                          sid=None,
+                          noise_scale=1.,
+                          length_scale=1.,
+                          noise_scale_w=1.,
+                          max_len=None,
+                          scope_shift=0):
+        x, m_p, logs_p, x_mask = self.enc_p(x, t, x_lengths)
+        if self.n_speakers > 0:
+            g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
+        else:
+            g = None
+
+        if self.use_sdp:
+            logw = self.dp(x,
+                           x_mask,
+                           g=g,
+                           reverse=True,
+                           noise_scale=noise_scale_w)
+        else:
+            logw = self.dp(x, x_mask, g=g)
+        w = torch.exp(logw) * x_mask * length_scale
+        w_ceil = torch.ceil(w)
+        y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
+        y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, None),
+                                 1).to(x_mask.dtype)
+        attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
+        attn = commons.generate_path(w_ceil, attn_mask)
+
+        m_p = torch.einsum('bctn, bdn -> bdt', attn, m_p)
+        logs_p = torch.einsum('bctn, bdn -> bdt', attn, logs_p)
+
+        z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
+        z = self.flow(z_p, y_mask, g=g, reverse=True)
+        z_spec, z_yin = torch.split(z,
+                                    self.inter_channels - self.yin_channels,
+                                    dim=1)
+        z_yin_crop = self.crop_scope([z_yin], scope_shift)[0]
+        z_crop = torch.cat([z_spec, z_yin_crop], dim=1)
+        decoder_inputs = z_crop * y_mask
+        return decoder_inputs, attn, y_mask, (z_crop, z, z_p, m_p, logs_p)
+
+    def infer_pre_lr(
+        self,
+        x,
+        t,
+        x_lengths,
+        sid=None,
+        length_scale=1,
+        noise_scale_w=1.,
+    ):
+        x, m_p, logs_p, x_mask = self.enc_p(x, t, x_lengths)
+        if self.n_speakers > 0:
+            g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
+        else:
+            g = None
+
+        if self.use_sdp:
+            logw = self.dp(x,
+                           x_mask,
+                           g=g,
+                           reverse=True,
+                           noise_scale=noise_scale_w)
+        else:
+            logw = self.dp(x, x_mask, g=g)
+        w = torch.exp(logw) * x_mask * length_scale
+        w_ceil = torch.ceil(w)
+        return w_ceil, x, m_p, logs_p, x_mask, g
+
+    def infer_lr(self, w_ceil, x, m_p, logs_p, x_mask):
+        y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
+        y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, None),
+                                 1).to(x_mask.dtype)
+        attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
+        attn = commons.generate_path(w_ceil, attn_mask)
+
+        m_p = torch.einsum('bctn, bdn -> bdt', attn, m_p)
+        logs_p = torch.einsum('bctn, bdn -> bdt', attn, logs_p)
+        return m_p, logs_p, y_mask
+
+    def infer_post_lr_pre_decoder(self,
+                                  m_p,
+                                  logs_p,
+                                  g,
+                                  y_mask,
+                                  noise_scale=1,
+                                  scope_shift=0):
+
+        z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
+        z = self.flow(z_p, y_mask, g=g, reverse=True)
+        z_spec, z_yin = torch.split(z,
+                                    self.inter_channels - self.yin_channels,
+                                    dim=1)
+
+        z_yin_crop = self.crop_scope([z_yin], scope_shift)[0]
+        z_crop = torch.cat([z_spec, z_yin_crop], dim=1)
+        decoder_inputs = z_crop * y_mask
+
+        return decoder_inputs, y_mask, (z_crop, z, z_p, m_p, logs_p)
+
+    def infer_decode_chunk(self, decoder_inputs, sid=None):
+        if self.n_speakers > 0:
+            g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
+        else:
+            g = None
+        return self.dec(decoder_inputs, g=g)
+
+

modules.py

@@ -0,0 +1,425 @@
+# from https://github.com/jaywalnut310/vits
+import math
+import torch
+from torch import nn
+from torch.nn import Conv1d
+from torch.nn import functional as F
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+import commons
+from commons import init_weights, get_padding
+from transforms import piecewise_rational_quadratic_transform
+
+
+LRELU_SLOPE = 0.1
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+class ConvReluNorm(nn.Module):
+    def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):
+        super().__init__()
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
+        assert n_layers > 1, "Number of layers should be larger than 0."
+
+        self.conv_layers = nn.ModuleList()
+        self.norm_layers = nn.ModuleList()
+        self.conv_layers.append(
+            nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size//2)
+        )
+        self.norm_layers.append(LayerNorm(hidden_channels))
+        self.relu_drop = nn.Sequential(
+            nn.ReLU(),
+            nn.Dropout(p_dropout))
+        for _ in range(n_layers-1):
+            self.conv_layers.append(nn.Conv1d(
+                hidden_channels, hidden_channels, kernel_size, padding=kernel_size//2)
+            )
+            self.norm_layers.append(LayerNorm(hidden_channels))
+        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
+
+    def forward(self, x, x_mask):
+        x_org = x
+        for i in range(self.n_layers):
+            x = self.conv_layers[i](x * x_mask)
+            x = self.norm_layers[i](x)
+            x = self.relu_drop(x)
+        x = x_org + self.proj(x)
+        return x * x_mask
+
+
+class DDSConv(nn.Module):
+    """Dialted and Depth-Separable Convolution"""
+    def __init__(self, channels, kernel_size, n_layers, p_dropout=0.):
+        super().__init__()
+        self.channels = channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
+
+        self.drop = nn.Dropout(p_dropout)
+        self.convs_sep = nn.ModuleList()
+        self.convs_1x1 = nn.ModuleList()
+        self.norms_1 = nn.ModuleList()
+        self.norms_2 = nn.ModuleList()
+        for i in range(n_layers):
+            dilation = kernel_size ** i
+            padding = (kernel_size * dilation - dilation) // 2
+            self.convs_sep.append(
+                nn.Conv1d(
+                    channels, 
+                    channels, 
+                    kernel_size,
+                    groups=channels, 
+                    dilation=dilation, 
+                    padding=padding
+                )
+            )
+            self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
+            self.norms_1.append(LayerNorm(channels))
+            self.norms_2.append(LayerNorm(channels))
+
+    def forward(self, x, x_mask, g=None):
+        if g is not None:
+            x = x + g
+        for i in range(self.n_layers):
+            y = self.convs_sep[i](x * x_mask)
+            y = self.norms_1[i](y)
+            y = F.gelu(y)
+            y = self.convs_1x1[i](y)
+            y = self.norms_2[i](y)
+            y = F.gelu(y)
+            y = self.drop(y)
+            x = x + y
+        return x * x_mask
+
+
+class WN(torch.nn.Module):
+    def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
+        super(WN, self).__init__()
+        assert(kernel_size % 2 == 1)
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size,
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+        self.p_dropout = p_dropout
+
+        self.in_layers = torch.nn.ModuleList()
+        self.res_skip_layers = torch.nn.ModuleList()
+        self.drop = nn.Dropout(p_dropout)
+
+        if gin_channels != 0:
+            cond_layer = torch.nn.Conv1d(gin_channels, 2*hidden_channels*n_layers, 1)
+            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')
+
+        for i in range(n_layers):
+            dilation = dilation_rate ** i
+            padding = int((kernel_size * dilation - dilation) / 2)
+            in_layer = torch.nn.Conv1d(
+                hidden_channels, 
+                2*hidden_channels, 
+                kernel_size,
+                dilation=dilation, 
+                padding=padding
+            )
+            in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
+            self.in_layers.append(in_layer)
+
+            # last one is not necessary
+            if i < n_layers - 1:
+                res_skip_channels = 2 * hidden_channels
+            else:
+                res_skip_channels = hidden_channels
+
+            res_skip_layer = torch.nn.Conv1d(
+                hidden_channels, res_skip_channels, 1
+            )
+            res_skip_layer = torch.nn.utils.weight_norm(
+                res_skip_layer, name='weight'
+            )
+            self.res_skip_layers.append(res_skip_layer)
+
+    def forward(self, x, x_mask, g=None, **kwargs):
+        output = torch.zeros_like(x)
+        n_channels_tensor = torch.IntTensor([self.hidden_channels])
+
+        if g is not None:
+            g = self.cond_layer(g)
+
+        for i in range(self.n_layers):
+            x_in = self.in_layers[i](x)
+            if g is not None:
+                cond_offset = i * 2 * self.hidden_channels
+                g_l = g[:, cond_offset:cond_offset+2*self.hidden_channels, :]
+            else:
+                g_l = torch.zeros_like(x_in)
+
+            acts = commons.fused_add_tanh_sigmoid_multiply(
+                x_in,
+                g_l,
+                n_channels_tensor
+            )
+            acts = self.drop(acts)
+
+            res_skip_acts = self.res_skip_layers[i](acts)
+            if i < self.n_layers - 1:
+                res_acts = res_skip_acts[:, :self.hidden_channels, :]
+                x = (x + res_acts) * x_mask
+                output = output + res_skip_acts[:, self.hidden_channels:, :]
+            else:
+                output = output + res_skip_acts
+        return output * x_mask
+
+    def remove_weight_norm(self):
+        if self.gin_channels != 0:
+            torch.nn.utils.remove_weight_norm(self.cond_layer)
+        for l in self.in_layers:
+            torch.nn.utils.remove_weight_norm(l)
+        for l in self.res_skip_layers:
+            torch.nn.utils.remove_weight_norm(l)
+
+
+class ResBlock1(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ResBlock1, self).__init__()
+        self.convs1 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                padding=get_padding(kernel_size, dilation[1]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
+                padding=get_padding(kernel_size, dilation[2])))
+        ])
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                padding=get_padding(kernel_size, 1)))
+        ])
+        self.convs2.apply(init_weights)
+
+    def forward(self, x, x_mask=None):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c2(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class ResBlock2(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
+        super(ResBlock2, self).__init__()
+        self.convs = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                padding=get_padding(kernel_size, dilation[1])))
+        ])
+        self.convs.apply(init_weights)
+
+    def forward(self, x, x_mask=None):
+        for c in self.convs:
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+class Log(nn.Module):
+    def forward(self, x, x_mask, reverse=False, **kwargs):
+        if not reverse:
+            y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
+            logdet = torch.sum(-y, [1, 2])
+            return y, logdet
+        else:
+            x = torch.exp(x) * x_mask
+            return x
+
+
+class Flip(nn.Module):
+    def forward(self, x, *args, reverse=False, **kwargs):
+        x = torch.flip(x, [1])
+        if not reverse:
+            logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
+            return x, logdet
+        else:
+            return x
+
+
+class ElementwiseAffine(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.channels = channels
+        self.m = nn.Parameter(torch.zeros(channels, 1))
+        self.logs = nn.Parameter(torch.zeros(channels, 1))
+
+    def forward(self, x, x_mask, reverse=False, **kwargs):
+        if not reverse:
+            y = self.m + torch.exp(self.logs) * x
+            y = y * x_mask
+            logdet = torch.sum(self.logs * x_mask, [1, 2])
+            return y, logdet
+        else:
+            x = (x - self.m) * torch.exp(-self.logs) * x_mask
+            return x
+
+
+class ResidualCouplingLayer(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        p_dropout=0,
+        gin_channels=0,
+        mean_only=False
+    ):
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.half_channels = channels // 2
+        self.mean_only = mean_only
+
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = WN(
+            hidden_channels, 
+            kernel_size, 
+            dilation_rate, 
+            n_layers, 
+            p_dropout=p_dropout, 
+            gin_channels=gin_channels
+        )
+        self.post = nn.Conv1d(
+            hidden_channels, self.half_channels * (2 - mean_only), 1
+        )
+        self.post.weight.data.zero_()
+        self.post.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, logs = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            logs = torch.zeros_like(m)
+
+        if not reverse:
+            x1 = m + x1 * torch.exp(logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(logs, [1, 2])
+            return x, logdet
+        else:
+            x1 = (x1 - m) * torch.exp(-logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            return x
+
+
+class ConvFlow(nn.Module):
+    def __init__(self, in_channels, filter_channels, kernel_size, n_layers, num_bins=10, tail_bound=5.0):
+        super().__init__()
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.num_bins = num_bins
+        self.tail_bound = tail_bound
+        self.half_channels = in_channels // 2
+
+        self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)
+        self.convs = DDSConv(
+            filter_channels, kernel_size, n_layers, p_dropout=0.
+        )
+        self.proj = nn.Conv1d(
+            filter_channels, self.half_channels * (num_bins * 3 - 1), 1
+        )
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels]*2, 1)
+        h = self.pre(x0)
+        h = self.convs(h, x_mask, g=g)
+        h = self.proj(h) * x_mask
+
+        b, c, t = x0.shape
+        # [b, cx?, t] -> [b, c, t, ?]
+        h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2)
+
+        unnormalized_widths = h[..., :self.num_bins] / \
+            math.sqrt(self.filter_channels)
+        unnormalized_heights = h[..., self.num_bins:2 * self.num_bins] / \
+            math.sqrt(self.filter_channels)
+        unnormalized_derivatives = h[..., 2 * self.num_bins:]
+
+        x1, logabsdet = piecewise_rational_quadratic_transform(
+            x1,
+            unnormalized_widths,
+            unnormalized_heights,
+            unnormalized_derivatives,
+            inverse=reverse,
+            tails='linear',
+            tail_bound=self.tail_bound
+        )
+
+        x = torch.cat([x0, x1], 1) * x_mask
+        logdet = torch.sum(logabsdet * x_mask, [1, 2])
+        if not reverse:
+            return x, logdet
+        else:
+            return x

pqmf.py

@@ -0,0 +1,136 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2020 Tomoki Hayashi
+#  MIT License (https://opensource.org/licenses/MIT)
+
+"""Pseudo QMF modules."""
+'''
+Copied from https://github.com/kan-bayashi/ParallelWaveGAN/blob/master/parallel_wavegan/layers/pqmf.py
+'''
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from scipy.signal import kaiser
+
+
+def design_prototype_filter(taps=62, cutoff_ratio=0.142, beta=9.0):
+    """Design prototype filter for PQMF.
+    This method is based on `A Kaiser window approach for the design of prototype
+    filters of cosine modulated filterbanks`_.
+    Args:
+        taps (int): The number of filter taps.
+        cutoff_ratio (float): Cut-off frequency ratio.
+        beta (float): Beta coefficient for kaiser window.
+    Returns:
+        ndarray: Impluse response of prototype filter (taps + 1,).
+    .. _`A Kaiser window approach for the design of prototype filters of cosine modulated filterbanks`:
+        https://ieeexplore.ieee.org/abstract/document/681427
+    """
+    # check the arguments are valid
+    assert taps % 2 == 0, "The number of taps mush be even number."
+    assert 0.0 < cutoff_ratio < 1.0, "Cutoff ratio must be > 0.0 and < 1.0."
+
+    # make initial filter
+    omega_c = np.pi * cutoff_ratio
+    with np.errstate(invalid="ignore"):
+        h_i = np.sin(omega_c * (np.arange(taps + 1) - 0.5 * taps)) / (
+            np.pi * (np.arange(taps + 1) - 0.5 * taps)
+        )
+    h_i[taps // 2] = np.cos(0) * cutoff_ratio  # fix nan due to indeterminate form
+
+    # apply kaiser window
+    w = kaiser(taps + 1, beta)
+    h = h_i * w
+
+    return h
+
+
+class PQMF(torch.nn.Module):
+    """PQMF module.
+    This module is based on `Near-perfect-reconstruction pseudo-QMF banks`_.
+    .. _`Near-perfect-reconstruction pseudo-QMF banks`:
+        https://ieeexplore.ieee.org/document/258122
+    """
+
+    def __init__(self, subbands=4, taps=62, cutoff_ratio=0.142, beta=9.0):
+        """Initilize PQMF module.
+        The cutoff_ratio and beta parameters are optimized for #subbands = 4.
+        See dicussion in https://github.com/kan-bayashi/ParallelWaveGAN/issues/195.
+        Args:
+            subbands (int): The number of subbands.
+            taps (int): The number of filter taps.
+            cutoff_ratio (float): Cut-off frequency ratio.
+            beta (float): Beta coefficient for kaiser window.
+        """
+        super(PQMF, self).__init__()
+
+        # build analysis & synthesis filter coefficients
+        h_proto = design_prototype_filter(taps, cutoff_ratio, beta)
+        h_analysis = np.zeros((subbands, len(h_proto)))
+        h_synthesis = np.zeros((subbands, len(h_proto)))
+        for k in range(subbands):
+            h_analysis[k] = (
+                2
+                * h_proto
+                * np.cos(
+                    (2 * k + 1)
+                    * (np.pi / (2 * subbands))
+                    * (np.arange(taps + 1) - (taps / 2))
+                    + (-1) ** k * np.pi / 4
+                )
+            )
+            h_synthesis[k] = (
+                2
+                * h_proto
+                * np.cos(
+                    (2 * k + 1)
+                    * (np.pi / (2 * subbands))
+                    * (np.arange(taps + 1) - (taps / 2))
+                    - (-1) ** k * np.pi / 4
+                )
+            )
+
+        # convert to tensor
+        analysis_filter = torch.Tensor(h_analysis).float().unsqueeze(1)
+        synthesis_filter = torch.Tensor(h_synthesis).float().unsqueeze(0)
+
+        # register coefficients as beffer
+        self.register_buffer("analysis_filter", analysis_filter)
+        self.register_buffer("synthesis_filter", synthesis_filter)
+
+        # filter for downsampling & upsampling
+        updown_filter = torch.zeros((subbands, subbands, subbands)).float()
+        for k in range(subbands):
+            updown_filter[k, k, 0] = 1.0
+        self.register_buffer("updown_filter", updown_filter)
+        self.subbands = subbands
+
+        # keep padding info
+        self.pad_fn = torch.nn.ConstantPad1d(taps // 2, 0.0)
+
+    def analysis(self, x):
+        """Analysis with PQMF.
+        Args:
+            x (Tensor): Input tensor (B, 1, T).
+        Returns:
+            Tensor: Output tensor (B, subbands, T // subbands).
+        """
+        x = F.conv1d(self.pad_fn(x), self.analysis_filter)
+        return F.conv1d(x, self.updown_filter, stride=self.subbands)
+
+    def synthesis(self, x):
+        """Synthesis with PQMF.
+        Args:
+            x (Tensor): Input tensor (B, subbands, T // subbands).
+        Returns:
+            Tensor: Output tensor (B, 1, T).
+        """
+        # NOTE(kan-bayashi): Power will be dreased so here multipy by # subbands.
+        #   Not sure this is the correct way, it is better to check again.
+        # TODO(kan-bayashi): Understand the reconstruction procedure
+        x = F.conv_transpose1d(
+            x, self.updown_filter * self.subbands, stride=self.subbands
+        )
+        return F.conv1d(self.pad_fn(x), self.synthesis_filter)

text/__init__.py

@@ -0,0 +1,68 @@
+""" from https://github.com/keithito/tacotron """
+import re
+from unicodedata import normalize
+
+from text.cleaners import collapse_whitespace
+from text.symbols import lang_to_dict, lang_to_dict_inverse
+
+
+def text_to_sequence(raw_text, lang):
+    '''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
+        lang: language of the input text
+    Returns:
+        List of integers corresponding to the symbols in the text
+    '''
+
+    _symbol_to_id = lang_to_dict(lang)
+    text = collapse_whitespace(raw_text)
+
+    if lang == 'ko_KR':    
+        text = normalize('NFKD', text)
+        sequence = [_symbol_to_id[symbol] for symbol in text]
+        tone = [0 for i in sequence]
+
+    elif lang == 'en_US':
+        _curly_re = re.compile(r'(.*?)\{(.+?)\}(.*)')
+        sequence = []
+
+        while len(text):
+            m = _curly_re.match(text)
+
+            if m is not None:
+                ar = m.group(1)
+                sequence += [_symbol_to_id[symbol] for symbol in ar]
+                ar = m.group(2)
+                sequence += [_symbol_to_id[symbol] for symbol in ar.split()]
+                text = m.group(3)
+            else:
+                sequence += [_symbol_to_id[symbol] for symbol in text]
+                break
+
+        tone = [0 for i in sequence]
+
+    else:
+        raise RuntimeError('Wrong type of lang')
+
+    assert len(sequence) == len(tone)
+    return sequence, tone
+
+
+def sequence_to_text(sequence, lang):
+    '''Converts a sequence of IDs back to a string'''
+    _id_to_symbol = lang_to_dict_inverse(lang)
+    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,89 @@
+""" from https://github.com/keithito/tacotron """
+
+'''
+Cleaners are transformations that run over the input text at both training and eval time.
+Cleaners can be selected by passing a comma-delimited list of cleaner names as the "cleaners"
+hyperparameter. Some cleaners are English-specific. You'll typically want to use:
+  1. "english_cleaners" for English text
+  2. "transliteration_cleaners" for non-English text that can be transliterated to ASCII using
+     the Unidecode library (https://pypi.python.org/pypi/Unidecode)
+  3. "basic_cleaners" if you do not want to transliterate (in this case, you should also update
+     the symbols in symbols.py to match your data).
+'''
+
+import re
+from unidecode import unidecode
+from unicodedata import normalize
+
+from .numbers import normalize_numbers
+
+
+# Regular expression matching whitespace:
+_whitespace_re = re.compile(r'\s+')
+
+# List of (regular expression, replacement) pairs for abbreviations:
+_abbreviations = [(re.compile('\\b%s\\.' % x[0], re.IGNORECASE), x[1]) for x in [
+    ('mrs', 'misess'),
+    ('mr', 'mister'),
+    ('dr', 'doctor'),
+    ('st', 'saint'),
+    ('co', 'company'),
+    ('jr', 'junior'),
+    ('maj', 'major'),
+    ('gen', 'general'),
+    ('drs', 'doctors'),
+    ('rev', 'reverend'),
+    ('lt', 'lieutenant'),
+    ('hon', 'honorable'),
+    ('sgt', 'sergeant'),
+    ('capt', 'captain'),
+    ('esq', 'esquire'),
+    ('ltd', 'limited'),
+    ('col', 'colonel'),
+    ('ft', 'fort'),
+]]
+
+_cht_norm = [(re.compile(r'[%s]' % x[0]), x[1]) for x in [
+    ('。．；', '.'),
+    ('，、', ', '),
+    ('？', '?'),
+    ('！', '!'),
+    ('─‧', '-'),
+    ('…', '...'),
+    ('《》「」『』〈〉（）', "'"),
+    ('：︰', ':'),
+    ('　', ' ')
+]]
+
+def expand_abbreviations(text):
+    for regex, replacement in _abbreviations:
+        text = re.sub(regex, replacement, text)
+    return text
+
+def expand_numbers(text):
+    return normalize_numbers(text)
+
+def lowercase(text):
+    return text.lower()
+
+def collapse_whitespace(text):
+    return re.sub(_whitespace_re, ' ', text)
+
+def convert_to_ascii(text):
+    return unidecode(text)
+
+def english_cleaners(text):
+    '''Pipeline for English text, including abbreviation expansion.'''
+    text = convert_to_ascii(text)
+    #text = lowercase(text)
+    text = expand_numbers(text)
+    text = expand_abbreviations(text)
+    text = collapse_whitespace(text)
+    return text
+
+def korean_cleaners(text):
+    '''Pipeline for Korean text, including collapses whitespace.'''
+    text = collapse_whitespace(text)
+    text = normalize('NFKD', text)
+    return text
+   

text/numbers.py

@@ -0,0 +1,71 @@
+""" from https://github.com/keithito/tacotron """
+
+import inflect
+import re
+
+
+_inflect = inflect.engine()
+_comma_number_re = re.compile(r'([0-9][0-9\,]+[0-9])')
+_decimal_number_re = re.compile(r'([0-9]+\.[0-9]+)')
+_pounds_re = re.compile(r'£([0-9\,]*[0-9]+)')
+_dollars_re = re.compile(r'\$([0-9\.\,]*[0-9]+)')
+_ordinal_re = re.compile(r'[0-9]+(st|nd|rd|th)')
+_number_re = re.compile(r'[0-9]+')
+
+
+def _remove_commas(m):
+    return m.group(1).replace(',', '')
+
+
+def _expand_decimal_point(m):
+    return m.group(1).replace('.', ' point ')
+
+
+def _expand_dollars(m):
+    match = m.group(1)
+    parts = match.split('.')
+    if len(parts) > 2:
+        return match + ' dollars'  # Unexpected format
+    dollars = int(parts[0]) if parts[0] else 0
+    cents = int(parts[1]) if len(parts) > 1 and parts[1] else 0
+    if dollars and cents:
+        dollar_unit = 'dollar' if dollars == 1 else 'dollars'
+        cent_unit = 'cent' if cents == 1 else 'cents'
+        return '%s %s, %s %s' % (dollars, dollar_unit, cents, cent_unit)
+    elif dollars:
+        dollar_unit = 'dollar' if dollars == 1 else 'dollars'
+        return '%s %s' % (dollars, dollar_unit)
+    elif cents:
+        cent_unit = 'cent' if cents == 1 else 'cents'
+        return '%s %s' % (cents, cent_unit)
+    else:
+        return 'zero dollars'
+
+
+def _expand_ordinal(m):
+    return _inflect.number_to_words(m.group(0))
+
+
+def _expand_number(m):
+    num = int(m.group(0))
+    if num > 1000 and num < 3000:
+        if num == 2000:
+            return 'two thousand'
+        elif num > 2000 and num < 2010:
+            return 'two thousand ' + _inflect.number_to_words(num % 100)
+        elif num % 100 == 0:
+            return _inflect.number_to_words(num // 100) + ' hundred'
+        else:
+            return _inflect.number_to_words(num, andword='', zero='oh', group=2).replace(', ', ' ')
+    else:
+        return _inflect.number_to_words(num, andword='')
+
+
+def normalize_numbers(text):
+    text = re.sub(_comma_number_re, _remove_commas, text)
+    text = re.sub(_pounds_re, r'\1 pounds', text)
+    text = re.sub(_dollars_re, _expand_dollars, text)
+    text = re.sub(_decimal_number_re, _expand_decimal_point, text)
+    text = re.sub(_ordinal_re, _expand_ordinal, text)
+    text = re.sub(_number_re, _expand_number, text)
+    return text

text/symbols.py

@@ -0,0 +1,38 @@
+_pad = '_'
+_punc = ";:,.!?¡¿—-…«»'“”~() "
+
+_jamo_leads = "".join([chr(_) for _ in range(0x1100, 0x1113)])
+_jamo_vowels = "".join([chr(_) for _ in range(0x1161, 0x1176)])
+_jamo_tails = "".join([chr(_) for _ in range(0x11A8, 0x11C3)])
+_kor_characters = _jamo_leads + _jamo_vowels + _jamo_tails
+
+_cmu_characters = [
+    'AA', 'AE', 'AH',
+    'AO', 'AW', 'AY',
+    'B', 'CH', 'D', 'DH', 'EH', 'ER', 'EY',
+    'F', 'G', 'HH', 'IH', 'IY',
+    'JH', 'K', 'L', 'M', 'N', 'NG', 'OW', 'OY',
+    'P', 'R', 'S', 'SH', 'T', 'TH', 'UH', 'UW',
+    'V', 'W', 'Y', 'Z', 'ZH'
+]
+
+
+lang_to_symbols = {
+    'common': [_pad] + list(_punc),
+    'ko_KR': list(_kor_characters), 
+    'en_US': _cmu_characters, 
+}
+
+def lang_to_dict(lang):
+    symbol_lang = lang_to_symbols['common'] + lang_to_symbols[lang]
+    dict_lang = {s: i for i, s in enumerate(symbol_lang)}
+    return dict_lang
+
+def lang_to_dict_inverse(lang):
+    symbol_lang = lang_to_symbols['common'] + lang_to_symbols[lang]
+    dict_lang = {i: s for i, s in enumerate(symbol_lang)}
+    return dict_lang
+
+def symbol_len(lang):
+    symbol_lang = lang_to_symbols['common'] + lang_to_symbols[lang]
+    return len(symbol_lang)

transforms.py

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

utils.py

@@ -0,0 +1,309 @@
+# from https://github.com/jaywalnut310/vits
+import os
+import sys
+import logging
+import subprocess
+import torch
+import numpy as np
+from omegaconf import OmegaConf
+from scipy.io.wavfile import read
+
+MATPLOTLIB_FLAG = False
+
+logging.basicConfig(
+    stream=sys.stdout, 
+    level=logging.INFO, 
+    format='[%(levelname)s|%(filename)s:%(lineno)s][%(asctime)s] >>> %(message)s'
+)
+logger = logging
+
+
+def load_checkpoint(checkpoint_path, rank=0, model_g=None, model_d=None, optim_g=None, optim_d=None):
+    assert os.path.isfile(checkpoint_path)
+    checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
+    iteration = checkpoint_dict['iteration']
+    learning_rate = checkpoint_dict['learning_rate']
+    config = checkpoint_dict['config']
+
+    if model_g is not None:
+        model_g, optim_g = load_model(
+            model_g, 
+            checkpoint_dict['model_g'],
+            optim_g,
+            checkpoint_dict['optimizer_g'])
+
+    if model_d is not None:
+        model_d, optim_d = load_model(
+            model_d, 
+            checkpoint_dict['model_d'],
+            optim_d,
+            checkpoint_dict['optimizer_d'])
+    if rank == 0:
+        logger.info(
+            "Loaded checkpoint '{}' (iteration {})".format(
+                checkpoint_path, 
+                iteration
+            )
+        )
+    return model_g, model_d, optim_g, optim_d, learning_rate, iteration, config
+    
+def load_checkpoint_diffsize(checkpoint_path, rank=0, model_g=None, model_d=None):
+    assert os.path.isfile(checkpoint_path)
+    checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
+    iteration = checkpoint_dict['iteration']
+    learning_rate = checkpoint_dict['learning_rate']
+    config = checkpoint_dict['config']
+
+    if model_g is not None:
+        model_g = load_model_diffsize(
+            model_g, 
+            checkpoint_dict['model_g'])
+    if model_d is not None:
+        model_d = load_model_diffsize(
+            model_d, 
+            checkpoint_dict['model_d'])
+    if rank == 0:
+        logger.info(
+            "Loaded checkpoint '{}' (iteration {})".format(
+                checkpoint_path, 
+                iteration
+            )
+        )
+    del checkpoint_dict
+    return model_g, model_d, learning_rate, iteration, config
+ 
+def load_model_diffsize(model, model_state_dict):
+    if hasattr(model, 'module'):
+        state_dict = model.module.state_dict()
+    else:
+        state_dict = model.state_dict()
+
+    for k, v in model_state_dict.items():
+        if k in state_dict and state_dict[k].size() == v.size():
+            state_dict[k] = v
+    
+    if hasattr(model, 'module'):
+        model.module.load_state_dict(state_dict, strict=False)
+    else:
+        model.load_state_dict(state_dict, strict=False)
+        
+    return model
+
+
+
+def load_model(model, model_state_dict, optim, optim_state_dict):
+    if optim is not None:
+        optim.load_state_dict(optim_state_dict)
+
+    if hasattr(model, 'module'):
+        state_dict = model.module.state_dict()
+    else:
+        state_dict = model.state_dict()
+
+    for k, v in model_state_dict.items():
+        if k in state_dict and state_dict[k].size() == v.size():
+            state_dict[k] = v
+    
+    if hasattr(model, 'module'):
+        model.module.load_state_dict(state_dict)
+    else:
+        model.load_state_dict(state_dict)
+        
+    return model, optim
+
+
+def save_checkpoint(net_g, optim_g, net_d, optim_d, hps, epoch, learning_rate, save_path):
+    
+    def get_state_dict(model):
+        if hasattr(model, 'module'):
+            state_dict = model.module.state_dict()
+        else:
+            state_dict = model.state_dict()
+        return state_dict
+    
+    torch.save({'model_g': get_state_dict(net_g),
+                'model_d': get_state_dict(net_d),
+                'optimizer_g': optim_g.state_dict(),
+                'optimizer_d': optim_d.state_dict(),
+                'config': str(hps),
+                'iteration': epoch,
+                'learning_rate': learning_rate}, save_path)
+
+
+def summarize(writer, global_step, scalars={}, histograms={}, images={}, audios={}, audio_sampling_rate=22050):
+    for k, v in scalars.items():
+        writer.add_scalar(k, v, global_step)
+    for k, v in histograms.items():
+        writer.add_histogram(k, v, global_step)
+    for k, v in images.items():
+        writer.add_image(k, v, global_step, dataformats='HWC')
+    for k, v in audios.items():
+        writer.add_audio(k, v, global_step, audio_sampling_rate)
+
+
+def plot_spectrogram_to_numpy(spectrogram):
+    global MATPLOTLIB_FLAG
+    if not MATPLOTLIB_FLAG:
+        import matplotlib
+        matplotlib.use("Agg")
+        MATPLOTLIB_FLAG = True
+        mpl_logger = logging.getLogger('matplotlib')
+        mpl_logger.setLevel(logging.WARNING)
+    import matplotlib.pylab as plt
+    import numpy as np
+
+    fig, ax = plt.subplots(figsize=(10, 2))
+    im = ax.imshow(spectrogram, aspect="auto", origin="lower",
+                   interpolation='none')
+    plt.colorbar(im, ax=ax)
+    plt.xlabel("Frames")
+    plt.ylabel("Channels")
+    plt.tight_layout()
+
+    fig.canvas.draw()
+    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    plt.close()
+    return data
+
+
+def plot_alignment_to_numpy(alignment, info=None):
+    global MATPLOTLIB_FLAG
+    if not MATPLOTLIB_FLAG:
+        import matplotlib
+        matplotlib.use("Agg")
+        MATPLOTLIB_FLAG = True
+        mpl_logger = logging.getLogger('matplotlib')
+        mpl_logger.setLevel(logging.WARNING)
+    import matplotlib.pylab as plt
+    import numpy as np
+
+    fig, ax = plt.subplots(figsize=(6, 4))
+    im = ax.imshow(alignment.transpose(), aspect='auto', origin='lower',
+                   interpolation='none')
+    fig.colorbar(im, ax=ax)
+    xlabel = 'Decoder timestep'
+    if info is not None:
+        xlabel += '\n\n' + info
+    plt.xlabel(xlabel)
+    plt.ylabel('Encoder timestep')
+    plt.tight_layout()
+
+    fig.canvas.draw()
+    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    plt.close()
+    return data
+
+
+def load_wav_to_torch(full_path):
+    sampling_rate, wav = read(full_path)
+    
+    if len(wav.shape) == 2:
+        wav = wav[:, 0]
+
+    if wav.dtype == np.int16:
+        wav = wav / 32768.0
+    elif wav.dtype == np.int32:
+        wav = wav / 2147483648.0
+    elif wav.dtype == np.uint8:
+        wav = (wav - 128) / 128.0
+    wav = wav.astype(np.float32)
+    return torch.FloatTensor(wav), sampling_rate
+
+
+def load_filepaths_and_text(filename, split="|"):
+    with open(filename, encoding='utf-8') as f:
+        filepaths_and_text = [line.strip().split(split) for line in f]
+    return filepaths_and_text
+
+
+def get_hparams(args, init=True):
+    config = OmegaConf.load(args.config)
+    hparams = HParams(**config)
+    model_dir = os.path.join(hparams.train.log_path, args.model)
+
+    if not os.path.exists(model_dir):
+        os.makedirs(model_dir)
+    hparams.model_name = args.model
+    hparams.model_dir = model_dir
+    config_save_path = os.path.join(model_dir, "config.yaml")
+
+    if init:
+        OmegaConf.save(config, config_save_path)
+
+    return hparams
+
+
+def get_hparams_from_file(config_path):
+    config = OmegaConf.load(config_path)
+    hparams = HParams(**config)
+    return hparams
+
+
+def check_git_hash(model_dir):
+    source_dir = os.path.dirname(os.path.realpath(__file__))
+    if not os.path.exists(os.path.join(source_dir, ".git")):
+        logger.warn("{} is not a git repository, therefore hash value comparison will be ignored.".format(
+            source_dir
+        ))
+        return
+
+    cur_hash = subprocess.getoutput("git rev-parse HEAD")
+
+    path = os.path.join(model_dir, "githash")
+    if os.path.exists(path):
+        saved_hash = open(path).read()
+        if saved_hash != cur_hash:
+            logger.warn("git hash values are different. {}(saved) != {}(current)".format(
+                saved_hash[:8], cur_hash[:8]))
+    else:
+        open(path, "w").write(cur_hash)
+
+
+def get_logger(model_dir, filename="train.log"):
+    global logger
+    logger = logging.getLogger(os.path.basename(model_dir))
+    logger.setLevel(logging.DEBUG)
+
+    formatter = logging.Formatter(
+        "%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")
+    if not os.path.exists(model_dir):
+        os.makedirs(model_dir)
+    h = logging.FileHandler(os.path.join(model_dir, filename))
+    h.setLevel(logging.DEBUG)
+    h.setFormatter(formatter)
+    logger.addHandler(h)
+    return logger
+
+
+class HParams():
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            if type(v) == dict:
+                v = HParams(**v)
+            self[k] = v
+
+    def keys(self):
+        return self.__dict__.keys()
+
+    def items(self):
+        return self.__dict__.items()
+
+    def values(self):
+        return self.__dict__.values()
+
+    def __len__(self):
+        return len(self.__dict__)
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __setitem__(self, key, value):
+        return setattr(self, key, value)
+
+    def __contains__(self, key):
+        return key in self.__dict__
+
+    def __repr__(self):
+        return self.__dict__.__repr__()

yin.py

@@ -0,0 +1,166 @@
+# remove np from https://github.com/dhchoi99/NANSY/blob/master/models/yin.py
+# adapted from https://github.com/patriceguyot/Yin
+# https://github.com/NVIDIA/mellotron/blob/master/yin.py
+
+import torch
+import torch.nn.functional as F
+from math import log2, ceil
+
+
+def differenceFunction(x, N, tau_max):
+    """
+    Compute difference function of data x. This corresponds to equation (6) in [1]
+    This solution is implemented directly with torch rfft.
+
+
+    :param x: audio data (Tensor)
+    :param N: length of data
+    :param tau_max: integration window size
+    :return: difference function
+    :rtype: list
+    """
+
+    #x = np.array(x, np.float64) #[B,T]
+    assert x.dim() == 2
+    b, w = x.shape
+    if w < tau_max:
+        x = F.pad(x, (tau_max - w - (tau_max - w) // 2, (tau_max - w) // 2),
+                  'constant',
+                  mode='reflect')
+    w = tau_max
+    #x_cumsum = np.concatenate((np.array([0.]), (x * x).cumsum()))
+    x_cumsum = torch.cat(
+        [torch.zeros([b, 1], device=x.device), (x * x).cumsum(dim=1)], dim=1)
+    size = w + tau_max
+    p2 = (size // 32).bit_length()
+    #p2 = ceil(log2(size+1 // 32))
+    nice_numbers = (16, 18, 20, 24, 25, 27, 30, 32)
+    size_pad = min(n * 2**p2 for n in nice_numbers if n * 2**p2 >= size)
+    fc = torch.fft.rfft(x, size_pad)  #[B,F]
+    conv = torch.fft.irfft(fc * fc.conj())[:, :tau_max]
+    return x_cumsum[:, w:w - tau_max:
+                    -1] + x_cumsum[:, w] - x_cumsum[:, :tau_max] - 2 * conv
+
+
+def differenceFunction_np(x, N, tau_max):
+    """
+    Compute difference function of data x. This corresponds to equation (6) in [1]
+    This solution is implemented directly with Numpy fft.
+
+
+    :param x: audio data
+    :param N: length of data
+    :param tau_max: integration window size
+    :return: difference function
+    :rtype: list
+    """
+
+    x = np.array(x, np.float64)
+    w = x.size
+    tau_max = min(tau_max, w)
+    x_cumsum = np.concatenate((np.array([0.]), (x * x).cumsum()))
+    size = w + tau_max
+    p2 = (size // 32).bit_length()
+    nice_numbers = (16, 18, 20, 24, 25, 27, 30, 32)
+    size_pad = min(x * 2**p2 for x in nice_numbers if x * 2**p2 >= size)
+    fc = np.fft.rfft(x, size_pad)
+    conv = np.fft.irfft(fc * fc.conjugate())[:tau_max]
+    return x_cumsum[w:w -
+                    tau_max:-1] + x_cumsum[w] - x_cumsum[:tau_max] - 2 * conv
+
+
+def cumulativeMeanNormalizedDifferenceFunction(df, N, eps=1e-8):
+    """
+    Compute cumulative mean normalized difference function (CMND).
+
+    This corresponds to equation (8) in [1]
+
+    :param df: Difference function
+    :param N: length of data
+    :return: cumulative mean normalized difference function
+    :rtype: list
+    """
+    #np.seterr(divide='ignore', invalid='ignore')
+    # scipy method, assert df>0 for all element
+    #   cmndf = df[1:] * np.asarray(list(range(1, N))) / (np.cumsum(df[1:]).astype(float) + eps)
+    B, _ = df.shape
+    cmndf = df[:,
+               1:] * torch.arange(1, N, device=df.device, dtype=df.dtype).view(
+                   1, -1) / (df[:, 1:].cumsum(dim=-1) + eps)
+    return torch.cat(
+        [torch.ones([B, 1], device=df.device, dtype=df.dtype), cmndf], dim=-1)
+
+
+def differenceFunctionTorch(xs: torch.Tensor, N, tau_max) -> torch.Tensor:
+    """pytorch backend batch-wise differenceFunction
+    has 1e-4 level error with input shape of (32, 22050*1.5)
+    Args:
+        xs:
+        N:
+        tau_max:
+
+    Returns:
+
+    """
+    xs = xs.double()
+    w = xs.shape[-1]
+    tau_max = min(tau_max, w)
+    zeros = torch.zeros((xs.shape[0], 1))
+    x_cumsum = torch.cat((torch.zeros((xs.shape[0], 1), device=xs.device),
+                          (xs * xs).cumsum(dim=-1, dtype=torch.double)),
+                         dim=-1)  # B x w
+    size = w + tau_max
+    p2 = (size // 32).bit_length()
+    nice_numbers = (16, 18, 20, 24, 25, 27, 30, 32)
+    size_pad = min(x * 2**p2 for x in nice_numbers if x * 2**p2 >= size)
+
+    fcs = torch.fft.rfft(xs, n=size_pad, dim=-1)
+    convs = torch.fft.irfft(fcs * fcs.conj())[:, :tau_max]
+    y1 = torch.flip(x_cumsum[:, w - tau_max + 1:w + 1], dims=[-1])
+    y = y1 + x_cumsum[:, w].unsqueeze(-1) - x_cumsum[:, :tau_max] - 2 * convs
+    return y
+
+
+def cumulativeMeanNormalizedDifferenceFunctionTorch(dfs: torch.Tensor,
+                                                    N,
+                                                    eps=1e-8) -> torch.Tensor:
+    arange = torch.arange(1, N, device=dfs.device, dtype=torch.float64)
+    cumsum = torch.cumsum(dfs[:, 1:], dim=-1,
+                          dtype=torch.float64).to(dfs.device)
+
+    cmndfs = dfs[:, 1:] * arange / (cumsum + eps)
+    cmndfs = torch.cat(
+        (torch.ones(cmndfs.shape[0], 1, device=dfs.device), cmndfs), dim=-1)
+    return cmndfs
+
+
+if __name__ == '__main__':
+    wav = torch.randn(32, int(22050 * 1.5)).cuda()
+    wav_numpy = wav.detach().cpu().numpy()
+    x = wav_numpy[0]
+
+    w_len = 2048
+    w_step = 256
+    tau_max = 2048
+    W = 2048
+
+    startFrames = list(range(0, x.shape[-1] - w_len, w_step))
+    startFrames = np.asarray(startFrames)
+    # times = startFrames / sr
+    frames = [x[..., t:t + W] for t in startFrames]
+    frames = np.asarray(frames)
+    frames_torch = torch.from_numpy(frames).cuda()
+
+    cmndfs0 = []
+    for idx, frame in enumerate(frames):
+        df = differenceFunction(frame, frame.shape[-1], tau_max)
+        cmndf = cumulativeMeanNormalizedDifferenceFunction(df, tau_max)
+        cmndfs0.append(cmndf)
+    cmndfs0 = np.asarray(cmndfs0)
+
+    dfs = differenceFunctionTorch(frames_torch, frames_torch.shape[-1],
+                                  tau_max)
+    cmndfs1 = cumulativeMeanNormalizedDifferenceFunctionTorch(
+        dfs, tau_max).detach().cpu().numpy()
+    print(cmndfs0.shape, cmndfs1.shape)
+    print(np.sum(np.abs(cmndfs0 - cmndfs1)))