Upload 10 files

app.py

@@ -0,0 +1,200 @@
+from huggingface_hub import snapshot_download
+from katsu import Katsu
+from models import build_model
+import gradio as gr
+import noisereduce as nr
+import numpy as np
+import os
+import phonemizer
+import random
+import torch
+import yaml
+
+random_texts = {}
+for lang in ['en', 'ja']:
+    with open(f'{lang}.txt', 'r') as r:
+        random_texts[lang] = [line.strip() for line in r]
+
+def get_random_text(voice):
+    if voice[0] == 'j':
+        lang = 'ja'
+    else:
+        lang = 'en'
+    return random.choice(random_texts[lang])
+
+def parens_to_angles(s):
+    return s.replace('(', '«').replace(')', '»')
+
+def normalize(text):
+    # TODO: Custom text normalization rules?
+    text = text.replace('Dr.', 'Doctor')
+    text = text.replace('Mr.', 'Mister')
+    text = text.replace('Ms.', 'Miss')
+    text = text.replace('Mrs.', 'Mrs')
+    return parens_to_angles(text)
+
+phonemizers = dict(
+    a=phonemizer.backend.EspeakBackend(language='en-us', preserve_punctuation=True, with_stress=True),
+    b=phonemizer.backend.EspeakBackend(language='en-gb', preserve_punctuation=True, with_stress=True),
+    j=Katsu()
+)
+
+def phonemize(text, voice):
+    lang = voice[0]
+    text = normalize(text)
+    ps = phonemizers[lang].phonemize([text])
+    ps = ps[0] if ps else ''
+    # TODO: Custom phonemization rules?
+    ps = parens_to_angles(ps)
+    # https://en.wiktionary.org/wiki/kokoro#English
+    ps = ps.replace('kəkˈoːɹoʊ', 'kˈoʊkəɹoʊ').replace('kəkˈɔːɹəʊ', 'kˈəʊkəɹəʊ')
+    ps = ''.join(filter(lambda p: p in VOCAB, ps))
+    return ps.strip()
+
+def length_to_mask(lengths):
+    mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
+    mask = torch.gt(mask+1, lengths.unsqueeze(1))
+    return mask
+
+def get_vocab():
+    _pad = "$"
+    _punctuation = ';:,.!?¡¿—…"«»“” '
+    _letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
+    _letters_ipa = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
+    symbols = [_pad] + list(_punctuation) + list(_letters) + list(_letters_ipa)
+    dicts = {}
+    for i in range(len((symbols))):
+        dicts[symbols[i]] = i
+    return dicts
+
+VOCAB = get_vocab()
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+snapshot = snapshot_download(repo_id='hexgrad/kokoro', allow_patterns=['*.pt', '*.pth', '*.yml'], use_auth_token=os.environ['TOKEN'])
+config = yaml.safe_load(open(os.path.join(snapshot, 'config.yml')))
+model = build_model(config['model_params'])
+_ = [model[key].eval() for key in model]
+_ = [model[key].to(device) for key in model]
+for key, state_dict in torch.load(os.path.join(snapshot, 'net.pth'), map_location='cpu', weights_only=True)['net'].items():
+    assert key in model, key
+    try:
+        model[key].load_state_dict(state_dict)
+    except:
+        state_dict = {k[7:]: v for k, v in state_dict.items()}
+        model[key].load_state_dict(state_dict, strict=False)
+
+CHOICES = {
+    '🇺🇸 🚺 American Female 0': 'af0',
+    '🇺🇸 🚺 Bella': 'af1',
+    '🇺🇸 🚺 Nicole': 'af2',
+    '🇺🇸 🚹 Michael': 'am0',
+    '🇺🇸 🚹 Adam': 'am1',
+    '🇬🇧 🚺 British Female 0': 'bf0',
+    '🇬🇧 🚺 British Female 1': 'bf1',
+    '🇬🇧 🚺 British Female 2': 'bf2',
+    '🇬🇧 🚹 British Male 0': 'bm0',
+    '🇬🇧 🚹 British Male 1': 'bm1',
+    '🇬🇧 🚹 British Male 2': 'bm2',
+    '🇬🇧 🚹 British Male 3': 'bm3',
+    '🇯🇵 🚺 Japanese Female 0': 'jf0',
+}
+VOICES = {k: torch.load(os.path.join(snapshot, 'voices', f'{k}.pt'), weights_only=True).to(device) for k in CHOICES.values()}
+
+np_log_99 = np.log(99)
+def s_curve(p):
+    if p <= 0:
+        return 0
+    elif p >= 1:
+        return 1
+    s = 1 / (1 + np.exp((1-p*2)*np_log_99))
+    s = (s-0.01) * 50/49
+    return s
+
+SAMPLE_RATE = 24000
+
+@torch.no_grad()
+def forward(text, voice, ps=None, speed=1.0, reduce_noise=0.5, opening_cut=5000, closing_cut=0, ease_in=3000, ease_out=0):
+    ps = ps or phonemize(text, voice)
+    tokens = [i for i in map(VOCAB.get, ps) if i is not None]
+    if not tokens:
+        return (None, '')
+    elif len(tokens) > 510:
+        tokens = tokens[:510]
+    ps = ''.join(next(k for k, v in VOCAB.items() if i == v) for i in tokens)
+    tokens = torch.LongTensor([[0, *tokens, 0]]).to(device)
+    input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)
+    text_mask = length_to_mask(input_lengths).to(device)
+    bert_dur = model.bert(tokens, attention_mask=(~text_mask).int())
+    d_en = model.bert_encoder(bert_dur).transpose(-1, -2)
+    ref_s = VOICES[voice]
+    s = ref_s[:, 128:]
+    d = model.predictor.text_encoder(d_en, s, input_lengths, text_mask)
+    x, _ = model.predictor.lstm(d)
+    duration = model.predictor.duration_proj(x)
+    duration = torch.sigmoid(duration).sum(axis=-1) / speed
+    pred_dur = torch.round(duration.squeeze()).clamp(min=1)
+    pred_aln_trg = torch.zeros(input_lengths, int(pred_dur.sum().data))
+    c_frame = 0
+    for i in range(pred_aln_trg.size(0)):
+        pred_aln_trg[i, c_frame:c_frame + int(pred_dur[i].data)] = 1
+        c_frame += int(pred_dur[i].data)
+    en = (d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0).to(device))
+    F0_pred, N_pred = model.predictor.F0Ntrain(en, s)
+    t_en = model.text_encoder(tokens, input_lengths, text_mask)
+    asr = (t_en @ pred_aln_trg.unsqueeze(0).to(device))
+    out = model.decoder(asr, F0_pred, N_pred, ref_s[:, :128])
+    out = out.squeeze().cpu().numpy()
+    if reduce_noise > 0:
+        out = nr.reduce_noise(y=out, sr=SAMPLE_RATE, prop_decrease=reduce_noise, n_fft=512)
+    opening_cut = max(0, int(opening_cut / speed))
+    if opening_cut > 0:
+        out[:opening_cut] = 0
+    closing_cut = max(0, int(closing_cut / speed))
+    if closing_cut > 0:
+        out = out[-closing_cut:] = 0
+    ease_in = min(int(ease_in / speed), len(out)//2 - opening_cut)
+    for i in range(ease_in):
+        out[i+opening_cut] *= s_curve(i / ease_in)
+    ease_out = min(int(ease_out / speed), len(out)//2 - closing_cut)
+    for i in range(ease_out):
+        out[-i-1-closing_cut] *= s_curve(i / ease_out)
+    return ((SAMPLE_RATE, out), ps)
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            text = gr.Textbox(label='Input Text')
+            voice = gr.Dropdown(list(CHOICES.items()), label='Voice')
+            with gr.Row():
+                random_btn = gr.Button('Random Text', variant='secondary')
+                generate_btn = gr.Button('Generate', variant='primary')
+            random_btn.click(get_random_text, inputs=[voice], outputs=[text])
+            with gr.Accordion('Input Phonemes', open=False):
+                in_ps = gr.Textbox(show_label=False, info='Override the input text with custom pronunciation. Leave this blank to use the input text instead.')
+                with gr.Row():
+                    clear_btn = gr.ClearButton(in_ps)
+                    phonemize_btn = gr.Button('Phonemize Input Text', variant='primary')
+            phonemize_btn.click(phonemize, inputs=[text, voice], outputs=[in_ps])
+        with gr.Column():
+            audio = gr.Audio(interactive=False, label='Output Audio')
+            with gr.Accordion('Tokens', open=True):
+                out_ps = gr.Textbox(interactive=False, show_label=False, info='Tokens used to generate the audio. Same as input phonemes if supplied, excluding unknown characters and truncated to 510 tokens.')
+    with gr.Accordion('Advanced Settings', open=False):
+        with gr.Row():
+            reduce_noise = gr.Slider(minimum=0, maximum=1, value=0.5, label='Reduce Noise', info='👻 Fix it in post: non-stationary noise reduction via spectral gating.')
+        with gr.Row():
+            speed = gr.Slider(minimum=0.5, maximum=2.0, value=1.0, step=0.1, label='Speed', info='⚡️ Adjust the speed of the audio. The trim settings below are also auto-scaled by speed.')
+        with gr.Row():
+            with gr.Column():
+                opening_cut = gr.Slider(minimum=0, maximum=24000, value=5000, step=1000, label='Opening Cut', info='✂️ Zero out this many samples at the start.')
+            with gr.Column():
+                closing_cut = gr.Slider(minimum=0, maximum=24000, value=0, step=1000, label='Closing Cut', info='✂️ Zero out this many samples at the end.')
+        with gr.Row():
+            with gr.Column():
+                ease_in = gr.Slider(minimum=0, maximum=24000, value=3000, step=1000, label='Ease In', info='🚀 Ease in for this many samples, after opening cut.')
+            with gr.Column():
+                ease_out = gr.Slider(minimum=0, maximum=24000, value=0, step=1000, label='Ease Out', info='📐 Ease out for this many samples, before closing cut.')
+    generate_btn.click(forward, inputs=[text, voice, in_ps, speed, reduce_noise, opening_cut, closing_cut, ease_in, ease_out], outputs=[audio, out_ps])
+
+if __name__ == '__main__':
+    demo.launch()

istftnet.py

@@ -0,0 +1,523 @@
+# https://github.com/yl4579/StyleTTS2/blob/main/Modules/istftnet.py
+from scipy.signal import get_window
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils import weight_norm, remove_weight_norm
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# https://github.com/yl4579/StyleTTS2/blob/main/Modules/utils.py
+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)
+
+LRELU_SLOPE = 0.1
+
+class AdaIN1d(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm1d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class AdaINResBlock1(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), style_dim=64):
+        super(AdaINResBlock1, self).__init__()
+        self.convs1 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                               padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                               padding=get_padding(kernel_size, dilation[1]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
+                               padding=get_padding(kernel_size, dilation[2])))
+        ])
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1)))
+        ])
+        self.convs2.apply(init_weights)
+        
+        self.adain1 = nn.ModuleList([
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+        ])
+        
+        self.adain2 = nn.ModuleList([
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+        ])
+        
+        self.alpha1 = nn.ParameterList([nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs1))])
+        self.alpha2 = nn.ParameterList([nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs2))])
+
+
+    def forward(self, x, s):
+        for c1, c2, n1, n2, a1, a2 in zip(self.convs1, self.convs2, self.adain1, self.adain2, self.alpha1, self.alpha2):
+            xt = n1(x, s)
+            xt = xt + (1 / a1) * (torch.sin(a1 * xt) ** 2)  # Snake1D
+            xt = c1(xt)
+            xt = n2(xt, s)
+            xt = xt + (1 / a2) * (torch.sin(a2 * xt) ** 2)  # Snake1D
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+            
+class TorchSTFT(torch.nn.Module):
+    def __init__(self, filter_length=800, hop_length=200, win_length=800, window='hann'):
+        super().__init__()
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.window = torch.from_numpy(get_window(window, win_length, fftbins=True).astype(np.float32))
+
+    def transform(self, input_data):
+        forward_transform = torch.stft(
+            input_data,
+            self.filter_length, self.hop_length, self.win_length, window=self.window.to(input_data.device),
+            return_complex=True)
+
+        return torch.abs(forward_transform), torch.angle(forward_transform)
+
+    def inverse(self, magnitude, phase):
+        inverse_transform = torch.istft(
+            magnitude * torch.exp(phase * 1j),
+            self.filter_length, self.hop_length, self.win_length, window=self.window.to(magnitude.device))
+
+        return inverse_transform.unsqueeze(-2)  # unsqueeze to stay consistent with conv_transpose1d implementation
+
+    def forward(self, input_data):
+        self.magnitude, self.phase = self.transform(input_data)
+        reconstruction = self.inverse(self.magnitude, self.phase)
+        return reconstruction
+    
+class SineGen(torch.nn.Module):
+    """ Definition of sine generator
+    SineGen(samp_rate, harmonic_num = 0,
+            sine_amp = 0.1, noise_std = 0.003,
+            voiced_threshold = 0,
+            flag_for_pulse=False)
+    samp_rate: sampling rate in Hz
+    harmonic_num: number of harmonic overtones (default 0)
+    sine_amp: amplitude of sine-wavefrom (default 0.1)
+    noise_std: std of Gaussian noise (default 0.003)
+    voiced_thoreshold: F0 threshold for U/V classification (default 0)
+    flag_for_pulse: this SinGen is used inside PulseGen (default False)
+    Note: when flag_for_pulse is True, the first time step of a voiced
+        segment is always sin(np.pi) or cos(0)
+    """
+
+    def __init__(self, samp_rate, upsample_scale, harmonic_num=0,
+                 sine_amp=0.1, noise_std=0.003,
+                 voiced_threshold=0,
+                 flag_for_pulse=False):
+        super(SineGen, self).__init__()
+        self.sine_amp = sine_amp
+        self.noise_std = noise_std
+        self.harmonic_num = harmonic_num
+        self.dim = self.harmonic_num + 1
+        self.sampling_rate = samp_rate
+        self.voiced_threshold = voiced_threshold
+        self.flag_for_pulse = flag_for_pulse
+        self.upsample_scale = upsample_scale
+
+    def _f02uv(self, f0):
+        # generate uv signal
+        uv = (f0 > self.voiced_threshold).type(torch.float32)
+        return uv
+
+    def _f02sine(self, f0_values):
+        """ f0_values: (batchsize, length, dim)
+            where dim indicates fundamental tone and overtones
+        """
+        # convert to F0 in rad. The interger part n can be ignored
+        # because 2 * np.pi * n doesn't affect phase
+        rad_values = (f0_values / self.sampling_rate) % 1
+
+        # initial phase noise (no noise for fundamental component)
+        rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], \
+                              device=f0_values.device)
+        rand_ini[:, 0] = 0
+        rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
+
+        # instantanouse phase sine[t] = sin(2*pi \sum_i=1 ^{t} rad)
+        if not self.flag_for_pulse:
+#             # for normal case
+
+#             # To prevent torch.cumsum numerical overflow,
+#             # it is necessary to add -1 whenever \sum_k=1^n rad_value_k > 1.
+#             # Buffer tmp_over_one_idx indicates the time step to add -1.
+#             # This will not change F0 of sine because (x-1) * 2*pi = x * 2*pi
+#             tmp_over_one = torch.cumsum(rad_values, 1) % 1
+#             tmp_over_one_idx = (padDiff(tmp_over_one)) < 0
+#             cumsum_shift = torch.zeros_like(rad_values)
+#             cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
+
+#             phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi
+            rad_values = torch.nn.functional.interpolate(rad_values.transpose(1, 2), 
+                                                         scale_factor=1/self.upsample_scale, 
+                                                         mode="linear").transpose(1, 2)
+    
+#             tmp_over_one = torch.cumsum(rad_values, 1) % 1
+#             tmp_over_one_idx = (padDiff(tmp_over_one)) < 0
+#             cumsum_shift = torch.zeros_like(rad_values)
+#             cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
+    
+            phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi
+            phase = torch.nn.functional.interpolate(phase.transpose(1, 2) * self.upsample_scale, 
+                                                    scale_factor=self.upsample_scale, mode="linear").transpose(1, 2)
+            sines = torch.sin(phase)
+            
+        else:
+            # If necessary, make sure that the first time step of every
+            # voiced segments is sin(pi) or cos(0)
+            # This is used for pulse-train generation
+
+            # identify the last time step in unvoiced segments
+            uv = self._f02uv(f0_values)
+            uv_1 = torch.roll(uv, shifts=-1, dims=1)
+            uv_1[:, -1, :] = 1
+            u_loc = (uv < 1) * (uv_1 > 0)
+
+            # get the instantanouse phase
+            tmp_cumsum = torch.cumsum(rad_values, dim=1)
+            # different batch needs to be processed differently
+            for idx in range(f0_values.shape[0]):
+                temp_sum = tmp_cumsum[idx, u_loc[idx, :, 0], :]
+                temp_sum[1:, :] = temp_sum[1:, :] - temp_sum[0:-1, :]
+                # stores the accumulation of i.phase within
+                # each voiced segments
+                tmp_cumsum[idx, :, :] = 0
+                tmp_cumsum[idx, u_loc[idx, :, 0], :] = temp_sum
+
+            # rad_values - tmp_cumsum: remove the accumulation of i.phase
+            # within the previous voiced segment.
+            i_phase = torch.cumsum(rad_values - tmp_cumsum, dim=1)
+
+            # get the sines
+            sines = torch.cos(i_phase * 2 * np.pi)
+        return sines
+
+    def forward(self, f0):
+        """ sine_tensor, uv = forward(f0)
+        input F0: tensor(batchsize=1, length, dim=1)
+                  f0 for unvoiced steps should be 0
+        output sine_tensor: tensor(batchsize=1, length, dim)
+        output uv: tensor(batchsize=1, length, 1)
+        """
+        f0_buf = torch.zeros(f0.shape[0], f0.shape[1], self.dim,
+                             device=f0.device)
+        # fundamental component
+        fn = torch.multiply(f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device))
+
+        # generate sine waveforms
+        sine_waves = self._f02sine(fn) * self.sine_amp
+
+        # generate uv signal
+        # uv = torch.ones(f0.shape)
+        # uv = uv * (f0 > self.voiced_threshold)
+        uv = self._f02uv(f0)
+
+        # noise: for unvoiced should be similar to sine_amp
+        #        std = self.sine_amp/3 -> max value ~ self.sine_amp
+        # .       for voiced regions is self.noise_std
+        noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
+        noise = noise_amp * torch.randn_like(sine_waves)
+
+        # first: set the unvoiced part to 0 by uv
+        # then: additive noise
+        sine_waves = sine_waves * uv + noise
+        return sine_waves, uv, noise
+
+
+class SourceModuleHnNSF(torch.nn.Module):
+    """ SourceModule for hn-nsf
+    SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshod=0)
+    sampling_rate: sampling_rate in Hz
+    harmonic_num: number of harmonic above F0 (default: 0)
+    sine_amp: amplitude of sine source signal (default: 0.1)
+    add_noise_std: std of additive Gaussian noise (default: 0.003)
+        note that amplitude of noise in unvoiced is decided
+        by sine_amp
+    voiced_threshold: threhold to set U/V given F0 (default: 0)
+    Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+    F0_sampled (batchsize, length, 1)
+    Sine_source (batchsize, length, 1)
+    noise_source (batchsize, length 1)
+    uv (batchsize, length, 1)
+    """
+
+    def __init__(self, sampling_rate, upsample_scale, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshod=0):
+        super(SourceModuleHnNSF, self).__init__()
+
+        self.sine_amp = sine_amp
+        self.noise_std = add_noise_std
+
+        # to produce sine waveforms
+        self.l_sin_gen = SineGen(sampling_rate, upsample_scale, harmonic_num,
+                                 sine_amp, add_noise_std, voiced_threshod)
+
+        # to merge source harmonics into a single excitation
+        self.l_linear = torch.nn.Linear(harmonic_num + 1, 1)
+        self.l_tanh = torch.nn.Tanh()
+
+    def forward(self, x):
+        """
+        Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+        F0_sampled (batchsize, length, 1)
+        Sine_source (batchsize, length, 1)
+        noise_source (batchsize, length 1)
+        """
+        # source for harmonic branch
+        with torch.no_grad():
+            sine_wavs, uv, _ = self.l_sin_gen(x)
+        sine_merge = self.l_tanh(self.l_linear(sine_wavs))
+
+        # source for noise branch, in the same shape as uv
+        noise = torch.randn_like(uv) * self.sine_amp / 3
+        return sine_merge, noise, uv
+def padDiff(x):
+    return F.pad(F.pad(x, (0,0,-1,1), 'constant', 0) - x, (0,0,0,-1), 'constant', 0)
+
+    
+class Generator(torch.nn.Module):
+    def __init__(self, style_dim, resblock_kernel_sizes, upsample_rates, upsample_initial_channel, resblock_dilation_sizes, upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size):
+        super(Generator, self).__init__()
+
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        resblock = AdaINResBlock1
+
+        self.m_source = SourceModuleHnNSF(
+                    sampling_rate=24000,
+                    upsample_scale=np.prod(upsample_rates) * gen_istft_hop_size,
+                    harmonic_num=8, voiced_threshod=10)
+        self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates) * gen_istft_hop_size)
+        self.noise_convs = nn.ModuleList()
+        self.noise_res = nn.ModuleList()
+        
+        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()
+        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, style_dim))
+                
+            c_cur = upsample_initial_channel // (2 ** (i + 1))
+            
+            if i + 1 < len(upsample_rates):  #
+                stride_f0 = np.prod(upsample_rates[i + 1:])
+                self.noise_convs.append(Conv1d(
+                    gen_istft_n_fft + 2, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=(stride_f0+1) // 2))
+                self.noise_res.append(resblock(c_cur, 7, [1,3,5], style_dim))
+            else:
+                self.noise_convs.append(Conv1d(gen_istft_n_fft + 2, c_cur, kernel_size=1))
+                self.noise_res.append(resblock(c_cur, 11, [1,3,5], style_dim))
+                
+                
+        self.post_n_fft = gen_istft_n_fft
+        self.conv_post = weight_norm(Conv1d(ch, self.post_n_fft + 2, 7, 1, padding=3))
+        self.ups.apply(init_weights)
+        self.conv_post.apply(init_weights)
+        self.reflection_pad = torch.nn.ReflectionPad1d((1, 0))
+        self.stft = TorchSTFT(filter_length=gen_istft_n_fft, hop_length=gen_istft_hop_size, win_length=gen_istft_n_fft)
+        
+        
+    def forward(self, x, s, f0):
+        with torch.no_grad():
+            f0 = self.f0_upsamp(f0[:, None]).transpose(1, 2)  # bs,n,t
+
+            har_source, noi_source, uv = self.m_source(f0)
+            har_source = har_source.transpose(1, 2).squeeze(1)
+            har_spec, har_phase = self.stft.transform(har_source)
+            har = torch.cat([har_spec, har_phase], dim=1)
+        
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            x_source = self.noise_convs[i](har)
+            x_source = self.noise_res[i](x_source, s)
+
+            x = self.ups[i](x)
+            if i == self.num_upsamples - 1:
+                x = self.reflection_pad(x)
+
+            x = x + x_source
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i*self.num_kernels+j](x, s)
+                else:
+                    xs += self.resblocks[i*self.num_kernels+j](x, s)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        spec = torch.exp(x[:,:self.post_n_fft // 2 + 1, :])
+        phase = torch.sin(x[:, self.post_n_fft // 2 + 1:, :])
+        return self.stft.inverse(spec, phase)
+    
+    def fw_phase(self, x, s):
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i*self.num_kernels+j](x, s)
+                else:
+                    xs += self.resblocks[i*self.num_kernels+j](x, s)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.reflection_pad(x)
+        x = self.conv_post(x)
+        spec = torch.exp(x[:,:self.post_n_fft // 2 + 1, :])
+        phase = torch.sin(x[:, self.post_n_fft // 2 + 1:, :])
+        return spec, phase
+
+    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()
+        remove_weight_norm(self.conv_pre)
+        remove_weight_norm(self.conv_post)
+
+        
+class AdainResBlk1d(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=64, actv=nn.LeakyReLU(0.2),
+                 upsample='none', dropout_p=0.0):
+        super().__init__()
+        self.actv = actv
+        self.upsample_type = upsample
+        self.upsample = UpSample1d(upsample)
+        self.learned_sc = dim_in != dim_out
+        self._build_weights(dim_in, dim_out, style_dim)
+        self.dropout = nn.Dropout(dropout_p)
+        
+        if upsample == 'none':
+            self.pool = nn.Identity()
+        else:
+            self.pool = weight_norm(nn.ConvTranspose1d(dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1, output_padding=1))
+        
+        
+    def _build_weights(self, dim_in, dim_out, style_dim):
+        self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
+        self.conv2 = weight_norm(nn.Conv1d(dim_out, dim_out, 3, 1, 1))
+        self.norm1 = AdaIN1d(style_dim, dim_in)
+        self.norm2 = AdaIN1d(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
+
+    def _shortcut(self, x):
+        x = self.upsample(x)
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        x = self.pool(x)
+        x = self.conv1(self.dropout(x))
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(self.dropout(x))
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / np.sqrt(2)
+        return out
+    
+class UpSample1d(nn.Module):
+    def __init__(self, layer_type):
+        super().__init__()
+        self.layer_type = layer_type
+
+    def forward(self, x):
+        if self.layer_type == 'none':
+            return x
+        else:
+            return F.interpolate(x, scale_factor=2, mode='nearest')
+
+class Decoder(nn.Module):
+    def __init__(self, dim_in=512, F0_channel=512, style_dim=64, dim_out=80, 
+                resblock_kernel_sizes = [3,7,11],
+                upsample_rates = [10, 6],
+                upsample_initial_channel=512,
+                resblock_dilation_sizes=[[1,3,5], [1,3,5], [1,3,5]],
+                upsample_kernel_sizes=[20, 12], 
+                gen_istft_n_fft=20, gen_istft_hop_size=5):
+        super().__init__()
+        
+        self.decode = nn.ModuleList()
+        
+        self.encode = AdainResBlk1d(dim_in + 2, 1024, style_dim)
+        
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 512, style_dim, upsample=True))
+
+        self.F0_conv = weight_norm(nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))
+        
+        self.N_conv = weight_norm(nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))
+        
+        self.asr_res = nn.Sequential(
+            weight_norm(nn.Conv1d(512, 64, kernel_size=1)),
+        )
+        
+        
+        self.generator = Generator(style_dim, resblock_kernel_sizes, upsample_rates, 
+                                   upsample_initial_channel, resblock_dilation_sizes, 
+                                   upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size)
+        
+    def forward(self, asr, F0_curve, N, s):
+        F0 = self.F0_conv(F0_curve.unsqueeze(1))
+        N = self.N_conv(N.unsqueeze(1))
+        
+        x = torch.cat([asr, F0, N], axis=1)
+        x = self.encode(x, s)
+        
+        asr_res = self.asr_res(asr)
+        
+        res = True
+        for block in self.decode:
+            if res:
+                x = torch.cat([x, asr_res, F0, N], axis=1)
+            x = block(x, s)
+            if block.upsample_type != "none":
+                res = False
+                
+        x = self.generator(x, s, F0_curve)
+        return x

katsu.py

@@ -0,0 +1,430 @@
+# https://github.com/polm/cutlet/blob/master/cutlet/cutlet.py
+from dataclasses import dataclass
+from fugashi import Tagger
+from num2kana import Convert
+import mojimoji
+import re
+import unicodedata
+
+HEPBURN = {
+chr(12449):'a',   #ァ
+chr(12450):'a',   #ア
+chr(12451):'i',   #ィ
+chr(12452):'i',   #イ
+chr(12453):'ɯ',   #ゥ
+chr(12454):'ɯ',   #ウ
+chr(12455):'e',   #ェ
+chr(12456):'e',   #エ
+chr(12457):'o',   #ォ
+chr(12458):'o',   #オ
+chr(12459):'ka',  #カ
+chr(12460):'ɡa',  #ガ
+chr(12461):'ki',  #キ
+chr(12462):'ɡi',  #ギ
+chr(12463):'kɯ',  #ク
+chr(12464):'ɡɯ',  #グ
+chr(12465):'ke',  #ケ
+chr(12466):'ɡe',  #ゲ
+chr(12467):'ko',  #コ
+chr(12468):'ɡo',  #ゴ
+chr(12469):'sa',  #サ
+chr(12470):'za',  #ザ
+chr(12471):'ɕi',  #シ
+chr(12472):'dʑi', #ジ
+chr(12473):'sɨ',  #ス
+chr(12474):'zɨ',  #ズ
+chr(12475):'se',  #セ
+chr(12476):'ze',  #ゼ
+chr(12477):'so',  #ソ
+chr(12478):'zo',  #ゾ
+chr(12479):'ta',  #タ
+chr(12480):'da',  #ダ
+chr(12481):'tɕi', #チ
+chr(12482):'dʑi', #ヂ
+# chr(12483)        #ッ
+chr(12484):'tsɨ', #ツ
+chr(12485):'zɨ',  #ヅ
+chr(12486):'te',  #テ
+chr(12487):'de',  #デ
+chr(12488):'to',  #ト
+chr(12489):'do',  #ド
+chr(12490):'na',  #ナ
+chr(12491):'ɲi',  #ニ
+chr(12492):'nɯ',  #ヌ
+chr(12493):'ne',  #ネ
+chr(12494):'no',  #ノ
+chr(12495):'ha',  #ハ
+chr(12496):'ba',  #バ
+chr(12497):'pa',  #パ
+chr(12498):'çi',  #ヒ
+chr(12499):'bi',  #ビ
+chr(12500):'pi',  #ピ
+chr(12501):'ɸɯ',  #フ
+chr(12502):'bɯ',  #ブ
+chr(12503):'pɯ',  #プ
+chr(12504):'he',  #ヘ
+chr(12505):'be',  #ベ
+chr(12506):'pe',  #ペ
+chr(12507):'ho',  #ホ
+chr(12508):'bo',  #ボ
+chr(12509):'po',  #ポ
+chr(12510):'ma',  #マ
+chr(12511):'mi',  #ミ
+chr(12512):'mɯ',  #ム
+chr(12513):'me',  #メ
+chr(12514):'mo',  #モ
+chr(12515):'ja',  #ャ
+chr(12516):'ja',  #ヤ
+chr(12517):'jɯ',  #ュ
+chr(12518):'jɯ',  #ユ
+chr(12519):'jo',  #ョ
+chr(12520):'jo',  #ヨ
+chr(12521):'ra',  #ラ
+chr(12522):'ri',  #リ
+chr(12523):'rɯ',  #ル
+chr(12524):'re',  #レ
+chr(12525):'ro',  #ロ
+chr(12526):'wa',  #ヮ
+chr(12527):'wa',  #ワ
+chr(12528):'i',   #ヰ
+chr(12529):'e',   #ヱ
+chr(12530):'o',   #ヲ
+# chr(12531)        #ン
+chr(12532):'vɯ',  #ヴ
+chr(12533):'ka',  #ヵ
+chr(12534):'ke',  #ヶ
+}
+assert len(HEPBURN) == 84 and all(i in {12483, 12531} or chr(i) in HEPBURN for i in range(12449, 12535))
+
+for k, v in list(HEPBURN.items()):
+    HEPBURN[chr(ord(k)-96)] = v
+assert len(HEPBURN) == 84*2
+
+HEPBURN.update({
+chr(12535):'va',  #ヷ
+chr(12536):'vi',  #ヸ
+chr(12537):'ve',  #ヹ
+chr(12538):'vo',  #ヺ
+})
+assert len(HEPBURN) == 84*2+4 and all(chr(i) in HEPBURN for i in range(12535, 12539))
+
+HEPBURN.update({
+chr(12784):'kɯ',  #ㇰ
+chr(12785):'ɕi',  #ㇱ
+chr(12786):'sɨ',  #ㇲ
+chr(12787):'to',  #ㇳ
+chr(12788):'nɯ',  #ㇴ
+chr(12789):'ha',  #ㇵ
+chr(12790):'çi',  #ㇶ
+chr(12791):'ɸɯ',  #ㇷ
+chr(12792):'he',  #ㇸ
+chr(12793):'ho',  #ㇹ
+chr(12794):'mɯ',  #ㇺ
+chr(12795):'ra',  #ㇻ
+chr(12796):'ri',  #ㇼ
+chr(12797):'rɯ',  #ㇽ
+chr(12798):'re',  #ㇾ
+chr(12799):'ro',  #ㇿ
+})
+assert len(HEPBURN) == 84*2+4+16 and all(chr(i) in HEPBURN for i in range(12784, 12800))
+
+HEPBURN.update({
+chr(12452)+chr(12455):'je',  #イェ
+chr(12454)+chr(12451):'wi',  #ウィ
+chr(12454)+chr(12455):'we',  #ウェ
+chr(12454)+chr(12457):'wo',  #ウォ
+chr(12461)+chr(12455):'kʲe', #キェ
+chr(12461)+chr(12515):'kʲa', #キャ
+chr(12461)+chr(12517):'kʲɨ', #キュ
+chr(12461)+chr(12519):'kʲo', #キョ
+chr(12462)+chr(12515):'ɡʲa', #ギャ
+chr(12462)+chr(12517):'ɡʲɨ', #ギュ
+chr(12462)+chr(12519):'ɡʲo', #ギョ
+chr(12463)+chr(12449):'kʷa', #クァ
+chr(12463)+chr(12451):'kʷi', #クィ
+chr(12463)+chr(12455):'kʷe', #クェ
+chr(12463)+chr(12457):'kʷo', #クォ
+chr(12464)+chr(12449):'ɡʷa', #グァ
+chr(12464)+chr(12451):'ɡʷi', #グィ
+chr(12464)+chr(12455):'ɡʷe', #グェ
+chr(12464)+chr(12457):'ɡʷo', #グォ
+chr(12471)+chr(12455):'ɕe',  #シェ
+chr(12471)+chr(12515):'ɕa',  #シャ
+chr(12471)+chr(12517):'ɕɨ',  #シュ
+chr(12471)+chr(12519):'ɕo',  #ショ
+chr(12472)+chr(12455):'dʑe', #ジェ
+chr(12472)+chr(12515):'dʑa', #ジャ
+chr(12472)+chr(12517):'dʑɨ', #ジュ
+chr(12472)+chr(12519):'dʑo', #ジョ
+chr(12481)+chr(12455):'tɕe', #チェ
+chr(12481)+chr(12515):'tɕa', #チャ
+chr(12481)+chr(12517):'tɕɨ', #チュ
+chr(12481)+chr(12519):'tɕo', #チョ
+chr(12482)+chr(12515):'dʑa', #ヂャ
+chr(12482)+chr(12517):'dʑɨ', #ヂュ
+chr(12482)+chr(12519):'dʑo', #ヂョ
+chr(12484)+chr(12449):'tsa', #ツァ
+chr(12484)+chr(12451):'tsi', #ツィ
+chr(12484)+chr(12455):'tse', #ツェ
+chr(12484)+chr(12457):'tso', #ツォ
+chr(12486)+chr(12451):'ti',  #ティ
+chr(12486)+chr(12517):'tʲɨ', #テュ
+chr(12487)+chr(12451):'di',  #ディ
+chr(12487)+chr(12517):'dʲɨ', #デュ
+chr(12488)+chr(12453):'tɯ',  #トゥ
+chr(12489)+chr(12453):'dɯ',  #ドゥ
+chr(12491)+chr(12455):'ɲe',  #ニェ
+chr(12491)+chr(12515):'ɲa',  #ニャ
+chr(12491)+chr(12517):'ɲɨ',  #ニュ
+chr(12491)+chr(12519):'ɲo',  #ニョ
+chr(12498)+chr(12455):'çe',  #ヒェ
+chr(12498)+chr(12515):'ça',  #ヒャ
+chr(12498)+chr(12517):'çɨ',  #ヒュ
+chr(12498)+chr(12519):'ço',  #ヒョ
+chr(12499)+chr(12515):'bʲa', #ビャ
+chr(12499)+chr(12517):'bʲɨ', #ビュ
+chr(12499)+chr(12519):'bʲo', #ビョ
+chr(12500)+chr(12515):'pʲa', #ピャ
+chr(12500)+chr(12517):'pʲɨ', #ピュ
+chr(12500)+chr(12519):'pʲo', #ピョ
+chr(12501)+chr(12449):'ɸa',  #ファ
+chr(12501)+chr(12451):'ɸi',  #フィ
+chr(12501)+chr(12455):'ɸe',  #フェ
+chr(12501)+chr(12457):'ɸo',  #フォ
+chr(12501)+chr(12517):'ɸʲɨ', #フュ
+chr(12501)+chr(12519):'ɸʲo', #フョ
+chr(12511)+chr(12515):'mʲa', #ミャ
+chr(12511)+chr(12517):'mʲɨ', #ミュ
+chr(12511)+chr(12519):'mʲo', #ミョ
+chr(12522)+chr(12515):'rʲa', #リャ
+chr(12522)+chr(12517):'rʲɨ', #リュ
+chr(12522)+chr(12519):'rʲo', #リョ
+chr(12532)+chr(12449):'va',  #ヴァ
+chr(12532)+chr(12451):'vi',  #ヴィ
+chr(12532)+chr(12455):'ve',  #ヴェ
+chr(12532)+chr(12457):'vo',  #ヴォ
+chr(12532)+chr(12517):'vʲɨ', #ヴュ
+chr(12532)+chr(12519):'vʲo', #ヴョ
+})
+assert len(HEPBURN) == 84*2+4+16+76
+
+for k, v in list(HEPBURN.items()):
+    if len(k) != 2:
+        continue
+    a, b = k
+    assert a in HEPBURN and b in HEPBURN, (a, b)
+    a = chr(ord(a)-96)
+    b = chr(ord(b)-96)
+    assert a in HEPBURN and b in HEPBURN, (a, b)
+    HEPBURN[a+b] = v
+assert len(HEPBURN) == 84*2+4+16+76*2
+
+HEPBURN.update({
+# symbols
+# 'ー': '-', # 長音符, only used when repeated
+'。': '.',
+'、': ',',
+'？': '?',
+'！': '!',
+'「': '"',
+'」': '"',
+'『': '"',
+'』': '"',
+'：': ':',
+'（': '(',
+'）': ')',
+'《': '(',
+'》': ')',
+'【': '[',
+'】': ']',
+'・': ' ',#'/',
+'，': ',',
+'～': '—',
+'〜': '—',
+'—': '—',
+'«': '«',
+'»': '»',
+
+# other
+'゚': '', # combining handakuten by itself, just discard
+'゙': '', # combining dakuten by itself
+})
+
+def add_dakuten(kk):
+    """Given a kana (single-character string), add a dakuten."""
+    try:
+        # ii = 'かきくけこさしすせそたちつてとはひふへほ'.index(kk)
+        ii = 'カキクケコサシスセソタチツテトハヒフヘホ'.index(kk)
+        return 'ガギグゲゴザジズゼゾダヂヅデドバビブベボ'[ii]
+        # return 'がぎぐげござじずぜぞだぢづでどばびぶべぼ'[ii]
+    except ValueError:
+        # this is normal if the input is nonsense
+        return None
+
+SUTEGANA = 'ャュョァィゥェォ' #'ゃゅょぁぃぅぇぉ'
+PUNCT = '\'".!?(),;:-'
+ODORI = '々〃ゝゞヽゞ'
+
+@dataclass
+class Token:
+    surface: str
+    space: bool # if a space should follow
+    def __str__(self):
+        sp = " " if self.space else ""
+        return f"{self.surface}{sp}"
+
+class Katsu:
+    def __init__(self):
+        """Create a Katsu object, which holds configuration as well as
+        tokenizer state.
+
+        Typical usage:
+
+        ```python
+        katsu = Katsu()
+        roma = katsu.romaji("カツカレーを食べた")
+        # "Cutlet curry wo tabeta"
+        ```
+        """
+        self.tagger = Tagger()
+        self.table = dict(HEPBURN) # make a copy so we can modify it
+        self.exceptions = {}
+
+    def romaji(self, text):
+        """Build a complete string from input text."""
+        if not text:
+            return ''
+        text = self._normalize_text(text)
+        words = self.tagger(text)
+        tokens = self._romaji_tokens(words)
+        out = ''.join([str(tok) for tok in tokens])
+        return re.sub(r'\s+', ' ', out.strip())
+
+    def phonemize(self, texts):
+        # espeak-ng API
+        return [self.romaji(text) for text in texts]
+
+    def _normalize_text(self, text):
+        """Given text, normalize variations in Japanese.
+
+        This specifically removes variations that are meaningless for romaji
+        conversion using the following steps:
+
+        - Unicode NFKC normalization
+        - Full-width Latin to half-width
+        - Half-width katakana to full-width
+        """
+        # perform unicode normalization
+        text = re.sub(r'[〜～](?=\d)', 'から', text) # wave dash range
+        text = unicodedata.normalize('NFKC', text)
+        # convert all full-width alphanum to half-width, since it can go out as-is
+        text = mojimoji.zen_to_han(text, kana=False)
+        # replace half-width katakana with full-width
+        text = mojimoji.han_to_zen(text, digit=False, ascii=False)
+        return ''.join([(' '+Convert(t)) if t.isdigit() else t for t in re.findall(r'\d+|\D+', text)])
+
+    def _romaji_tokens(self, words):
+        """Build a list of tokens from input nodes."""
+        out = []
+        for wi, word in enumerate(words):
+            po = out[-1] if out else None
+            pw = words[wi - 1] if wi > 0 else None
+            nw = words[wi + 1] if wi < len(words) - 1 else None
+            roma = self._romaji_word(word)
+            tok = Token(roma, False)
+            # handle punctuation with atypical spacing
+            surface = word.surface#['orig']
+            if surface in '「『' or roma in '([':
+                if po:
+                    po.space = True
+            elif surface in '」』' or roma in ']).,?!:':
+                if po:
+                    po.space = False
+                tok.space = True
+            elif roma == ' ':
+                tok.space = False
+            else:
+                tok.space = True
+            out.append(tok)
+        # remove any leftover sokuon
+        for tok in out:
+            tok.surface = tok.surface.replace(chr(12483), '')
+        return out
+
+    def _romaji_word(self, word):
+        """Return the romaji for a single word (node)."""
+        surface = word.surface#['orig']
+        if surface in self.exceptions:
+            return self.exceptions[surface]
+        assert not surface.isdigit(), surface
+        if surface.isascii():
+            return surface
+        kana = word.feature.pron or word.feature.kana or surface
+        if word.is_unk:
+            if word.char_type == 7: # katakana
+                pass
+            elif word.char_type == 3: # symbol
+                return ''.join(map(lambda c: self.table.get(c, c), surface))
+            else:
+                return '' # TODO: silently fail
+        out = ''
+        for ki, char in enumerate(kana):
+            nk = kana[ki + 1] if ki < len(kana) - 1 else None
+            pk = kana[ki - 1] if ki > 0 else None
+            out += self._get_single_mapping(pk, char, nk)
+        return out
+
+    def _get_single_mapping(self, pk, kk, nk):
+        """Given a single kana and its neighbors, return the mapped romaji."""
+        # handle odoriji
+        # NOTE: This is very rarely useful at present because odoriji are not
+        # left in readings for dictionary words, and we can't follow kana
+        # across word boundaries.
+        if kk in ODORI:
+            if kk in 'ゝヽ':
+                if pk: return pk
+                else: return '' # invalid but be nice
+            if kk in 'ゞヾ': # repeat with voicing
+                if not pk: return ''
+                vv = add_dakuten(pk)
+                if vv: return self.table[vv]
+                else: return ''
+            # remaining are 々 for kanji and 〃 for symbols, but we can't
+            # infer their span reliably (or handle rendaku)
+            return ''
+        # handle digraphs
+        if pk and (pk + kk) in self.table:
+            return self.table[pk + kk]
+        if nk and (kk + nk) in self.table:
+            return ''
+        if nk and nk in SUTEGANA:
+            if kk == 'ッ': return '' # never valid, just ignore
+            return self.table[kk][:-1] + self.table[nk]
+        if kk in SUTEGANA:
+            return ''
+        if kk == 'ー': # 長音符
+            return 'ː'
+        if ord(kk) in {12387, 12483}: # っ or ッ
+            tnk = self.table.get(nk)
+            if tnk and tnk[0] in 'bdɸɡhçijkmnɲoprstɯvwz':
+                return tnk[0]
+            return kk
+        if ord(kk) in {12435, 12531}: # ん or ン
+            # https://en.wikipedia.org/wiki/N_(kana)
+            # m before m,p,b
+            # ŋ before k,g
+            # ɲ before ɲ,tɕ,dʑ
+            # n before n,t,d,r,z
+            # ɴ otherwise
+            tnk = self.table.get(nk)
+            if tnk:
+                if tnk[0] in 'mpb':
+                    return 'm'
+                elif tnk[0] in 'kɡ':
+                    return 'ŋ'
+                elif any(tnk.startswith(p) for p in ('ɲ','tɕ','dʑ')):
+                    return 'ɲ'
+                elif tnk[0] in 'ntdrz':
+                    return 'n'
+            return 'ɴ'
+        return self.table.get(kk, '')

models.py

@@ -0,0 +1,571 @@
+# https://github.com/yl4579/StyleTTS2/blob/main/models.py
+from istftnet import Decoder
+from munch import Munch
+from plbert import load_plbert
+from torch.nn.utils import weight_norm, spectral_norm
+import numpy as np
+import os
+import os.path as osp
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class LearnedDownSample(nn.Module):
+    def __init__(self, layer_type, dim_in):
+        super().__init__()
+        self.layer_type = layer_type
+
+        if self.layer_type == 'none':
+            self.conv = nn.Identity()
+        elif self.layer_type == 'timepreserve':
+            self.conv = spectral_norm(nn.Conv2d(dim_in, dim_in, kernel_size=(3, 1), stride=(2, 1), groups=dim_in, padding=(1, 0)))
+        elif self.layer_type == 'half':
+            self.conv = spectral_norm(nn.Conv2d(dim_in, dim_in, kernel_size=(3, 3), stride=(2, 2), groups=dim_in, padding=1))
+        else:
+            raise RuntimeError('Got unexpected donwsampletype %s, expected is [none, timepreserve, half]' % self.layer_type)
+            
+    def forward(self, x):
+        return self.conv(x)
+
+class LearnedUpSample(nn.Module):
+    def __init__(self, layer_type, dim_in):
+        super().__init__()
+        self.layer_type = layer_type
+        
+        if self.layer_type == 'none':
+            self.conv = nn.Identity()
+        elif self.layer_type == 'timepreserve':
+            self.conv = nn.ConvTranspose2d(dim_in, dim_in, kernel_size=(3, 1), stride=(2, 1), groups=dim_in, output_padding=(1, 0), padding=(1, 0))
+        elif self.layer_type == 'half':
+            self.conv = nn.ConvTranspose2d(dim_in, dim_in, kernel_size=(3, 3), stride=(2, 2), groups=dim_in, output_padding=1, padding=1)
+        else:
+            raise RuntimeError('Got unexpected upsampletype %s, expected is [none, timepreserve, half]' % self.layer_type)
+
+
+    def forward(self, x):
+        return self.conv(x)
+
+class DownSample(nn.Module):
+    def __init__(self, layer_type):
+        super().__init__()
+        self.layer_type = layer_type
+
+    def forward(self, x):
+        if self.layer_type == 'none':
+            return x
+        elif self.layer_type == 'timepreserve':
+            return F.avg_pool2d(x, (2, 1))
+        elif self.layer_type == 'half':
+            if x.shape[-1] % 2 != 0:
+                x = torch.cat([x, x[..., -1].unsqueeze(-1)], dim=-1)
+            return F.avg_pool2d(x, 2)
+        else:
+            raise RuntimeError('Got unexpected donwsampletype %s, expected is [none, timepreserve, half]' % self.layer_type)
+
+
+class UpSample(nn.Module):
+    def __init__(self, layer_type):
+        super().__init__()
+        self.layer_type = layer_type
+
+    def forward(self, x):
+        if self.layer_type == 'none':
+            return x
+        elif self.layer_type == 'timepreserve':
+            return F.interpolate(x, scale_factor=(2, 1), mode='nearest')
+        elif self.layer_type == 'half':
+            return F.interpolate(x, scale_factor=2, mode='nearest')
+        else:
+            raise RuntimeError('Got unexpected upsampletype %s, expected is [none, timepreserve, half]' % self.layer_type)
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize=False, downsample='none'):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = DownSample(downsample)
+        self.downsample_res = LearnedDownSample(downsample, dim_in)
+        self.learned_sc = dim_in != dim_out
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = spectral_norm(nn.Conv2d(dim_in, dim_in, 3, 1, 1))
+        self.conv2 = spectral_norm(nn.Conv2d(dim_in, dim_out, 3, 1, 1))
+        if self.normalize:
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        if self.learned_sc:
+            self.conv1x1 = spectral_norm(nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False))
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = self.downsample(x)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        x = self.downsample_res(x)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x / np.sqrt(2)  # unit variance
+
+class LinearNorm(torch.nn.Module):
+    def __init__(self, in_dim, out_dim, bias=True, w_init_gain='linear'):
+        super(LinearNorm, self).__init__()
+        self.linear_layer = torch.nn.Linear(in_dim, out_dim, bias=bias)
+
+        torch.nn.init.xavier_uniform_(
+            self.linear_layer.weight,
+            gain=torch.nn.init.calculate_gain(w_init_gain))
+
+    def forward(self, x):
+        return self.linear_layer(x)
+
+class Discriminator2d(nn.Module):
+    def __init__(self, dim_in=48, num_domains=1, max_conv_dim=384, repeat_num=4):
+        super().__init__()
+        blocks = []
+        blocks += [spectral_norm(nn.Conv2d(1, dim_in, 3, 1, 1))]
+
+        for lid in range(repeat_num):
+            dim_out = min(dim_in*2, max_conv_dim)
+            blocks += [ResBlk(dim_in, dim_out, downsample='half')]
+            dim_in = dim_out
+
+        blocks += [nn.LeakyReLU(0.2)]
+        blocks += [spectral_norm(nn.Conv2d(dim_out, dim_out, 5, 1, 0))]
+        blocks += [nn.LeakyReLU(0.2)]
+        blocks += [nn.AdaptiveAvgPool2d(1)]
+        blocks += [spectral_norm(nn.Conv2d(dim_out, num_domains, 1, 1, 0))]
+        self.main = nn.Sequential(*blocks)
+
+    def get_feature(self, x):
+        features = []
+        for l in self.main:
+            x = l(x)
+            features.append(x) 
+        out = features[-1]
+        out = out.view(out.size(0), -1)  # (batch, num_domains)
+        return out, features
+
+    def forward(self, x):
+        out, features = self.get_feature(x)
+        out = out.squeeze()  # (batch)
+        return out, features
+
+class ResBlk1d(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize=False, downsample='none', dropout_p=0.2):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample_type = downsample
+        self.learned_sc = dim_in != dim_out
+        self._build_weights(dim_in, dim_out)
+        self.dropout_p = dropout_p
+        
+        if self.downsample_type == 'none':
+            self.pool = nn.Identity()
+        else:
+            self.pool = weight_norm(nn.Conv1d(dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1))
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_in, 3, 1, 1))
+        self.conv2 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
+        if self.normalize:
+            self.norm1 = nn.InstanceNorm1d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm1d(dim_in, affine=True)
+        if self.learned_sc:
+            self.conv1x1 = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
+
+    def downsample(self, x):
+        if self.downsample_type == 'none':
+            return x
+        else:
+            if x.shape[-1] % 2 != 0:
+                x = torch.cat([x, x[..., -1].unsqueeze(-1)], dim=-1)
+            return F.avg_pool1d(x, 2)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        x = self.downsample(x)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = F.dropout(x, p=self.dropout_p, training=self.training)
+        
+        x = self.conv1(x)
+        x = self.pool(x)
+        if self.normalize:
+            x = self.norm2(x)
+            
+        x = self.actv(x)
+        x = F.dropout(x, p=self.dropout_p, training=self.training)
+        
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x / np.sqrt(2)  # unit variance
+
+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 TextEncoder(nn.Module):
+    def __init__(self, channels, kernel_size, depth, n_symbols, actv=nn.LeakyReLU(0.2)):
+        super().__init__()
+        self.embedding = nn.Embedding(n_symbols, channels)
+
+        padding = (kernel_size - 1) // 2
+        self.cnn = nn.ModuleList()
+        for _ in range(depth):
+            self.cnn.append(nn.Sequential(
+                weight_norm(nn.Conv1d(channels, channels, kernel_size=kernel_size, padding=padding)),
+                LayerNorm(channels),
+                actv,
+                nn.Dropout(0.2),
+            ))
+        # self.cnn = nn.Sequential(*self.cnn)
+
+        self.lstm = nn.LSTM(channels, channels//2, 1, batch_first=True, bidirectional=True)
+
+    def forward(self, x, input_lengths, m):
+        x = self.embedding(x)  # [B, T, emb]
+        x = x.transpose(1, 2)  # [B, emb, T]
+        m = m.to(input_lengths.device).unsqueeze(1)
+        x.masked_fill_(m, 0.0)
+        
+        for c in self.cnn:
+            x = c(x)
+            x.masked_fill_(m, 0.0)
+            
+        x = x.transpose(1, 2)  # [B, T, chn]
+
+        input_lengths = input_lengths.cpu().numpy()
+        x = nn.utils.rnn.pack_padded_sequence(
+            x, input_lengths, batch_first=True, enforce_sorted=False)
+
+        self.lstm.flatten_parameters()
+        x, _ = self.lstm(x)
+        x, _ = nn.utils.rnn.pad_packed_sequence(
+            x, batch_first=True)
+                
+        x = x.transpose(-1, -2)
+        x_pad = torch.zeros([x.shape[0], x.shape[1], m.shape[-1]])
+
+        x_pad[:, :, :x.shape[-1]] = x
+        x = x_pad.to(x.device)
+        
+        x.masked_fill_(m, 0.0)
+        
+        return x
+
+    def inference(self, x):
+        x = self.embedding(x)
+        x = x.transpose(1, 2)
+        x = self.cnn(x)
+        x = x.transpose(1, 2)
+        self.lstm.flatten_parameters()
+        x, _ = self.lstm(x)
+        return x
+    
+    def length_to_mask(self, lengths):
+        mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
+        mask = torch.gt(mask+1, lengths.unsqueeze(1))
+        return mask
+
+
+
+class AdaIN1d(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm1d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class UpSample1d(nn.Module):
+    def __init__(self, layer_type):
+        super().__init__()
+        self.layer_type = layer_type
+
+    def forward(self, x):
+        if self.layer_type == 'none':
+            return x
+        else:
+            return F.interpolate(x, scale_factor=2, mode='nearest')
+
+class AdainResBlk1d(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=64, actv=nn.LeakyReLU(0.2),
+                 upsample='none', dropout_p=0.0):
+        super().__init__()
+        self.actv = actv
+        self.upsample_type = upsample
+        self.upsample = UpSample1d(upsample)
+        self.learned_sc = dim_in != dim_out
+        self._build_weights(dim_in, dim_out, style_dim)
+        self.dropout = nn.Dropout(dropout_p)
+        
+        if upsample == 'none':
+            self.pool = nn.Identity()
+        else:
+            self.pool = weight_norm(nn.ConvTranspose1d(dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1, output_padding=1))
+        
+        
+    def _build_weights(self, dim_in, dim_out, style_dim):
+        self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
+        self.conv2 = weight_norm(nn.Conv1d(dim_out, dim_out, 3, 1, 1))
+        self.norm1 = AdaIN1d(style_dim, dim_in)
+        self.norm2 = AdaIN1d(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
+
+    def _shortcut(self, x):
+        x = self.upsample(x)
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        x = self.pool(x)
+        x = self.conv1(self.dropout(x))
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(self.dropout(x))
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / np.sqrt(2)
+        return out
+    
+class AdaLayerNorm(nn.Module):
+    def __init__(self, style_dim, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.fc = nn.Linear(style_dim, channels*2)
+
+    def forward(self, x, s):
+        x = x.transpose(-1, -2)
+        x = x.transpose(1, -1)
+                
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        gamma, beta = gamma.transpose(1, -1), beta.transpose(1, -1)
+        
+        
+        x = F.layer_norm(x, (self.channels,), eps=self.eps)
+        x = (1 + gamma) * x + beta
+        return x.transpose(1, -1).transpose(-1, -2)
+
+class ProsodyPredictor(nn.Module):
+
+    def __init__(self, style_dim, d_hid, nlayers, max_dur=50, dropout=0.1):
+        super().__init__() 
+        
+        self.text_encoder = DurationEncoder(sty_dim=style_dim, 
+                                            d_model=d_hid,
+                                            nlayers=nlayers, 
+                                            dropout=dropout)
+
+        self.lstm = nn.LSTM(d_hid + style_dim, d_hid // 2, 1, batch_first=True, bidirectional=True)
+        self.duration_proj = LinearNorm(d_hid, max_dur)
+        
+        self.shared = nn.LSTM(d_hid + style_dim, d_hid // 2, 1, batch_first=True, bidirectional=True)
+        self.F0 = nn.ModuleList()
+        self.F0.append(AdainResBlk1d(d_hid, d_hid, style_dim, dropout_p=dropout))
+        self.F0.append(AdainResBlk1d(d_hid, d_hid // 2, style_dim, upsample=True, dropout_p=dropout))
+        self.F0.append(AdainResBlk1d(d_hid // 2, d_hid // 2, style_dim, dropout_p=dropout))
+
+        self.N = nn.ModuleList()
+        self.N.append(AdainResBlk1d(d_hid, d_hid, style_dim, dropout_p=dropout))
+        self.N.append(AdainResBlk1d(d_hid, d_hid // 2, style_dim, upsample=True, dropout_p=dropout))
+        self.N.append(AdainResBlk1d(d_hid // 2, d_hid // 2, style_dim, dropout_p=dropout))
+        
+        self.F0_proj = nn.Conv1d(d_hid // 2, 1, 1, 1, 0)
+        self.N_proj = nn.Conv1d(d_hid // 2, 1, 1, 1, 0)
+
+
+    def forward(self, texts, style, text_lengths, alignment, m):
+        d = self.text_encoder(texts, style, text_lengths, m)
+        
+        batch_size = d.shape[0]
+        text_size = d.shape[1]
+        
+        # predict duration
+        input_lengths = text_lengths.cpu().numpy()
+        x = nn.utils.rnn.pack_padded_sequence(
+            d, input_lengths, batch_first=True, enforce_sorted=False)
+        
+        m = m.to(text_lengths.device).unsqueeze(1)
+        
+        self.lstm.flatten_parameters()
+        x, _ = self.lstm(x)
+        x, _ = nn.utils.rnn.pad_packed_sequence(
+            x, batch_first=True)
+        
+        x_pad = torch.zeros([x.shape[0], m.shape[-1], x.shape[-1]])
+
+        x_pad[:, :x.shape[1], :] = x
+        x = x_pad.to(x.device)
+                
+        duration = self.duration_proj(nn.functional.dropout(x, 0.5, training=self.training))
+        
+        en = (d.transpose(-1, -2) @ alignment)
+
+        return duration.squeeze(-1), en
+    
+    def F0Ntrain(self, x, s):
+        x, _ = self.shared(x.transpose(-1, -2))
+        
+        F0 = x.transpose(-1, -2)
+        for block in self.F0:
+            F0 = block(F0, s)
+        F0 = self.F0_proj(F0)
+
+        N = x.transpose(-1, -2)
+        for block in self.N:
+            N = block(N, s)
+        N = self.N_proj(N)
+        
+        return F0.squeeze(1), N.squeeze(1)
+    
+    def length_to_mask(self, lengths):
+        mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
+        mask = torch.gt(mask+1, lengths.unsqueeze(1))
+        return mask
+    
+class DurationEncoder(nn.Module):
+
+    def __init__(self, sty_dim, d_model, nlayers, dropout=0.1):
+        super().__init__()
+        self.lstms = nn.ModuleList()
+        for _ in range(nlayers):
+            self.lstms.append(nn.LSTM(d_model + sty_dim, 
+                                 d_model // 2, 
+                                 num_layers=1, 
+                                 batch_first=True, 
+                                 bidirectional=True, 
+                                 dropout=dropout))
+            self.lstms.append(AdaLayerNorm(sty_dim, d_model))
+        
+        
+        self.dropout = dropout
+        self.d_model = d_model
+        self.sty_dim = sty_dim
+
+    def forward(self, x, style, text_lengths, m):
+        masks = m.to(text_lengths.device)
+        
+        x = x.permute(2, 0, 1)
+        s = style.expand(x.shape[0], x.shape[1], -1)
+        x = torch.cat([x, s], axis=-1)
+        x.masked_fill_(masks.unsqueeze(-1).transpose(0, 1), 0.0)
+                
+        x = x.transpose(0, 1)
+        input_lengths = text_lengths.cpu().numpy()
+        x = x.transpose(-1, -2)
+        
+        for block in self.lstms:
+            if isinstance(block, AdaLayerNorm):
+                x = block(x.transpose(-1, -2), style).transpose(-1, -2)
+                x = torch.cat([x, s.permute(1, -1, 0)], axis=1)
+                x.masked_fill_(masks.unsqueeze(-1).transpose(-1, -2), 0.0)
+            else:
+                x = x.transpose(-1, -2)
+                x = nn.utils.rnn.pack_padded_sequence(
+                    x, input_lengths, batch_first=True, enforce_sorted=False)
+                block.flatten_parameters()
+                x, _ = block(x)
+                x, _ = nn.utils.rnn.pad_packed_sequence(
+                    x, batch_first=True)
+                x = F.dropout(x, p=self.dropout, training=self.training)
+                x = x.transpose(-1, -2)
+                
+                x_pad = torch.zeros([x.shape[0], x.shape[1], m.shape[-1]])
+
+                x_pad[:, :, :x.shape[-1]] = x
+                x = x_pad.to(x.device)
+        
+        return x.transpose(-1, -2)
+    
+    def inference(self, x, style):
+        x = self.embedding(x.transpose(-1, -2)) * np.sqrt(self.d_model)
+        style = style.expand(x.shape[0], x.shape[1], -1)
+        x = torch.cat([x, style], axis=-1)
+        src = self.pos_encoder(x)
+        output = self.transformer_encoder(src).transpose(0, 1)
+        return output
+    
+    def length_to_mask(self, lengths):
+        mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
+        mask = torch.gt(mask+1, lengths.unsqueeze(1))
+        return mask
+
+# https://github.com/yl4579/StyleTTS2/blob/main/utils.py
+def recursive_munch(d):
+    if isinstance(d, dict):
+        return Munch((k, recursive_munch(v)) for k, v in d.items())
+    elif isinstance(d, list):
+        return [recursive_munch(v) for v in d]
+    else:
+        return d
+
+def build_model(args):
+    args = recursive_munch(args)
+    assert args.decoder.type == 'istftnet', 'Decoder type unknown'
+    decoder = Decoder(dim_in=args.hidden_dim, style_dim=args.style_dim, dim_out=args.n_mels,
+            resblock_kernel_sizes = args.decoder.resblock_kernel_sizes,
+            upsample_rates = args.decoder.upsample_rates,
+            upsample_initial_channel=args.decoder.upsample_initial_channel,
+            resblock_dilation_sizes=args.decoder.resblock_dilation_sizes,
+            upsample_kernel_sizes=args.decoder.upsample_kernel_sizes,
+            gen_istft_n_fft=args.decoder.gen_istft_n_fft, gen_istft_hop_size=args.decoder.gen_istft_hop_size)
+    text_encoder = TextEncoder(channels=args.hidden_dim, kernel_size=5, depth=args.n_layer, n_symbols=args.n_token)
+    predictor = ProsodyPredictor(style_dim=args.style_dim, d_hid=args.hidden_dim, nlayers=args.n_layer, max_dur=args.max_dur, dropout=args.dropout)
+    bert = load_plbert()
+    return Munch(
+        bert=bert,
+        bert_encoder=nn.Linear(bert.config.hidden_size, args.hidden_dim),
+        predictor=predictor,
+        decoder=decoder,
+        text_encoder=text_encoder,
+    )

num2kana.py

@@ -0,0 +1,317 @@
+# https://github.com/Greatdane/Convert-Numbers-to-Japanese/blob/master/Convert-Numbers-to-Japanese.py
+# Japanese Number Converter
+# - Currently using length functions - possible to use Recursive Functions? Difficult with the many exceptions
+# - Works up to 9 figures (999999999)
+
+romaji_dict = {".": "ten", "0": "zero", "1": "ichi", "2": "ni", "3": "san", "4": "yon", "5": "go", "6": "roku", "7": "nana",
+                "8": "hachi", "9": "kyuu", "10": "juu", "100": "hyaku", "1000": "sen", "10000": "man", "100000000": "oku",
+                "300": "sanbyaku", "600": "roppyaku", "800": "happyaku", "3000": "sanzen", "8000":"hassen", "01000": "issen"}
+
+kanji_dict = {".": "点", "0": "零", "1": "一", "2": "二", "3": "三", "4": "四", "5": "五", "6": "六", "7": "七",
+                "8": "八", "9": "九", "10": "十", "100": "百", "1000": "千", "10000": "万", "100000000": "億",
+              "300": "三百", "600": "六百", "800": "八百", "3000": "三千", "8000":"八千", "01000": "一千"}
+
+hiragana_dict = {".": "てん", "0": "ゼロ", "1": "いち", "2": "に", "3": "さん", "4": "よん", "5": "ご", "6": "ろく", "7": "なな",
+                "8": "はち", "9": "きゅう", "10": "じゅう", "100": "ひゃく", "1000": "せん", "10000": "まん", "100000000": "おく",
+                 "300": "さんびゃく", "600": "ろっぴゃく", "800": "はっぴゃく", "3000": "さんぜん", "8000":"はっせん", "01000": "いっせん" }
+
+key_dict = {"kanji" : kanji_dict, "hiragana" : hiragana_dict, "romaji": romaji_dict}
+
+def len_one(convert_num,requested_dict):
+    # Returns single digit conversion, 0-9
+    return requested_dict[convert_num]
+
+def len_two(convert_num,requested_dict):
+    # Returns the conversion, when number is of length two (10-99)
+    if convert_num[0] == "0": #if 0 is first, return len_one
+        return len_one(convert_num[1],requested_dict)
+    if convert_num == "10":
+        return requested_dict["10"] # Exception, if number is 10, simple return 10
+    if convert_num[0] == "1": # When first number is 1, use ten plus second number
+        return requested_dict["10"] + " " + len_one(convert_num[1],requested_dict)
+    elif convert_num[1] == "0": # If ending number is zero, give first number plus 10
+        return len_one(convert_num[0],requested_dict) + " " + requested_dict["10"]
+    else:
+        num_list = []
+        for x in convert_num:
+            num_list.append(requested_dict[x])
+        num_list.insert(1, requested_dict["10"])
+        # Convert to a string (from a list)
+        output = ""
+        for y in num_list:
+            output += y + " "
+        output = output[:len(output) - 1]  # take off the space
+        return output
+
+def len_three(convert_num,requested_dict):
+    # Returns the conversion, when number is of length three (100-999)
+    num_list = []
+    if convert_num[0] == "1":
+        num_list.append(requested_dict["100"])
+    elif convert_num[0] == "3":
+        num_list.append(requested_dict["300"])
+    elif convert_num[0] == "6":
+        num_list.append(requested_dict["600"])
+    elif convert_num[0] == "8":
+        num_list.append(requested_dict["800"])
+    else:
+        num_list.append(requested_dict[convert_num[0]])
+        num_list.append(requested_dict["100"])
+    if convert_num[1:] == "00" and len(convert_num) == 3:
+        pass
+    else:
+        if convert_num[1] == "0":
+            num_list.append(requested_dict[convert_num[2]])
+        else:
+            num_list.append(len_two(convert_num[1:], requested_dict))
+    output = ""  
+    for y in num_list:
+         output += y + " "
+    output = output[:len(output) - 1]
+    return output
+
+def len_four(convert_num,requested_dict, stand_alone):
+    # Returns the conversion, when number is of length four (1000-9999)
+    num_list = []
+    # First, check for zeros (and get deal with them)
+    if convert_num == "0000":
+        return ""
+    while convert_num[0] == "0":
+        convert_num = convert_num[1:]
+    if len(convert_num) == 1:
+        return len_one(convert_num,requested_dict)
+    elif len(convert_num) == 2:
+        return len_two(convert_num,requested_dict)
+    elif len(convert_num) == 3:
+        return len_three(convert_num,requested_dict)
+    # If no zeros, do the calculation
+    else:
+        # Have to handle 1000, depending on if its a standalone 1000-9999 or included in a larger number
+        if convert_num[0] == "1" and stand_alone:
+            num_list.append(requested_dict["1000"])
+        elif convert_num[0] == "1":
+            num_list.append(requested_dict["01000"])
+        elif convert_num[0] == "3":
+            num_list.append(requested_dict["3000"])
+        elif convert_num[0] == "8":
+            num_list.append(requested_dict["8000"])
+        else:
+            num_list.append(requested_dict[convert_num[0]])
+            num_list.append(requested_dict["1000"])
+        if convert_num[1:] == "000" and len(convert_num) == 4:
+            pass
+        else:
+            if convert_num[1] == "0":
+                num_list.append(len_two(convert_num[2:],requested_dict))
+            else:
+                num_list.append(len_three(convert_num[1:],requested_dict))
+        output = ""
+        for y in num_list:
+             output += y + " "
+        output = output[:len(output) - 1]
+        return output
+
+
+def len_x(convert_num,requested_dict):
+    #Returns everything else.. (up to 9 digits)
+    num_list = []
+    if len(convert_num[0:-4]) == 1:
+        num_list.append(requested_dict[convert_num[0:-4]])
+        num_list.append(requested_dict["10000"])
+    elif len(convert_num[0:-4]) == 2:
+        num_list.append(len_two(convert_num[0:2],requested_dict))
+        num_list.append(requested_dict["10000"])
+    elif len(convert_num[0:-4]) == 3:
+        num_list.append(len_three(convert_num[0:3],requested_dict))
+        num_list.append(requested_dict["10000"])
+    elif len(convert_num[0:-4]) == 4:
+        num_list.append(len_four(convert_num[0:4],requested_dict, False))
+        num_list.append(requested_dict["10000"])
+    elif len(convert_num[0:-4]) == 5:
+        num_list.append(requested_dict[convert_num[0]])
+        num_list.append(requested_dict["100000000"])
+        num_list.append(len_four(convert_num[1:5],requested_dict, False))
+        if convert_num[1:5] == "0000":
+            pass
+        else:
+            num_list.append(requested_dict["10000"])
+    else:
+        assert False, "Not yet implemented, please choose a lower number."
+    num_list.append(len_four(convert_num[-4:],requested_dict, False))
+    output = ""
+    for y in num_list:
+         output += y + " "
+    output = output[:len(output) - 1]       
+    return output
+
+def remove_spaces(convert_result):
+    # Remove spaces in Hirigana and Kanji results
+    correction = ""
+    for x in convert_result:
+        if x == " ":
+            pass
+        else:
+         correction += x
+    return correction
+
+def do_convert(convert_num,requested_dict):
+    #Check lengths and convert accordingly
+    if len(convert_num) == 1:
+        return(len_one(convert_num,requested_dict))
+    elif len(convert_num) == 2:
+        return(len_two(convert_num,requested_dict))
+    elif len(convert_num) == 3:
+        return(len_three(convert_num,requested_dict))
+    elif len(convert_num) == 4:
+        return(len_four(convert_num,requested_dict, True))
+    else:
+        return(len_x(convert_num,requested_dict))
+
+def split_Point(split_num,dict_choice):
+    # Used if a decmial point is in the string.
+    split_num = split_num.split(".")
+    split_num_a = split_num[0]
+    split_num_b = split_num[1]
+    split_num_b_end = " "
+    for x in split_num_b:
+        split_num_b_end += len_one(x,key_dict[dict_choice]) + " "
+    # To account for small exception of small tsu when ending in jyuu in hiragana/romaji
+    if split_num_a[-1] == "0" and split_num_a[-2] != "0" and dict_choice == "hiragana":
+        small_Tsu = Convert(split_num_a,dict_choice)
+        small_Tsu = small_Tsu[0:-1] + "っ"
+        return small_Tsu + key_dict[dict_choice]["."] + split_num_b_end
+    if split_num_a[-1] == "0" and split_num_a[-2] != "0" and dict_choice == "romaji":
+        small_Tsu = Convert(split_num_a,dict_choice)
+        small_Tsu = small_Tsu[0:-1] + "t"
+        return small_Tsu + key_dict[dict_choice]["."] + split_num_b_end
+
+    return Convert(split_num_a,dict_choice) + " " + key_dict[dict_choice]["."] + split_num_b_end
+
+
+def do_kanji_convert(convert_num):
+    # Converts kanji to arabic number
+
+    if convert_num == "零":
+        return 0
+
+    # First, needs to check for MAN 万 and OKU 億 kanji, as need to handle differently, splitting up the numbers at these intervals.
+    # key tells us whether we need to add or multiply the numbers, then we create a list of numbers in an order we need to add/multiply
+    key = []
+    numberList = []
+    y = ""
+    for x in convert_num:
+        if x == "万" or x == "億":
+            numberList.append(y)
+            key.append("times")
+            numberList.append(x)
+            key.append("plus")
+            y = ""
+        else:
+            y += x
+    if y != "":
+        numberList.append(y)
+
+    numberListConverted = []
+    baseNumber = ["一", "二", "三", "四", "五", "六", "七", "八", "九"]
+    linkNumber = ["十", "百", "千", "万", "億"]
+
+    # Converts the kanji number list to arabic numbers, using the 'base number' and 'link number' list above. For a link number, we would need to
+    # link with a base number
+    for noX in numberList:
+        count = len(noX)
+        result = 0
+        skip = 1
+        for x in reversed(noX):
+            addTo = 0
+            skip -= 1
+            count = count - 1
+            if skip == 1:
+                continue
+            if x in baseNumber:
+                for y, z in kanji_dict.items():
+                    if z == x:
+                        result += int(y)
+            elif x in linkNumber:
+                if noX[count - 1] in baseNumber and count > 0:
+                    for y, z in kanji_dict.items():
+                        if z == noX[count - 1]:
+                            tempNo = int(y)
+                            for y, z in kanji_dict.items():
+                                if z == x:
+                                    addTo += tempNo * int(y)
+                                    result += addTo
+                                    skip = 2
+                else:
+                    for y, z in kanji_dict.items():
+                        if z == x:
+                            result += int(y)
+        numberListConverted.append(int(result))
+
+    result = numberListConverted[0]
+    y = 0
+
+    # Iterate over the converted list, and either multiply/add as instructed in key list
+    for x in range(1,len(numberListConverted)):
+        if key[y] == "plus":
+            try:
+                if key[y+1] == "times":
+                    result = result + numberListConverted[x] * numberListConverted[x+1]
+                    y += 1
+                else:
+                    result += numberListConverted[x]
+            except IndexError:
+                result += numberListConverted[-1]
+                break
+        else:
+            result = result * numberListConverted[x]
+        y += 1
+
+    return result
+
+def Convert(convert_num, dict_choice='hiragana'):
+    # Input formatting
+    convert_num = str(convert_num)
+    convert_num = convert_num.replace(',','')
+    dict_choice = dict_choice.lower()
+
+    # If all is selected as dict_choice, return as a list
+    if dict_choice == "all":
+        result_list = []
+        for x in "kanji", "hiragana", "romaji":
+            result_list.append(Convert(convert_num,x))
+        return result_list
+
+    dictionary = key_dict[dict_choice]
+    
+    # Exit if length is greater than current limit
+    if len(convert_num) > 9:
+        return("Number length too long, choose less than 10 digits")
+
+    # Remove any leading zeroes
+    while convert_num[0] == "0" and len(convert_num) > 1:
+        convert_num = convert_num[1:]
+
+    # Check for decimal places
+    if "." in convert_num:
+        result = split_Point(convert_num,dict_choice)
+    else:
+        result = do_convert(convert_num, dictionary)
+
+    # Remove spaces and return result
+    if key_dict[dict_choice] == romaji_dict:
+        pass
+    else:
+        result = remove_spaces(result)
+    return result
+
+def ConvertKanji(convert_num):
+    if convert_num[0] in kanji_dict.values():
+        # Check to see if 点 (point) is in the input, and handle by splitting at 点, before and after is handled separately
+        if "点" in convert_num:
+            point = convert_num.find("点")
+            endNumber = ""
+            for x in convert_num[point+1:]:
+                endNumber += list(kanji_dict.keys())[list(kanji_dict.values()).index(x)]
+            return(str(do_kanji_convert(convert_num[0:point])) + "." + endNumber)
+        else:
+            return(str(do_kanji_convert(convert_num)))

packages.txt

@@ -0,0 +1 @@
+espeak-ng

plbert.py

@@ -0,0 +1,15 @@
+# https://github.com/yl4579/StyleTTS2/blob/main/Utils/PLBERT/util.py
+from transformers import AlbertConfig, AlbertModel
+
+class CustomAlbert(AlbertModel):
+    def forward(self, *args, **kwargs):
+        # Call the original forward method
+        outputs = super().forward(*args, **kwargs)
+        # Only return the last_hidden_state
+        return outputs.last_hidden_state
+
+def load_plbert():
+    plbert_config = {'vocab_size': 178, 'hidden_size': 768, 'num_attention_heads': 12, 'intermediate_size': 2048, 'max_position_embeddings': 512, 'num_hidden_layers': 12, 'dropout': 0.1}
+    albert_base_configuration = AlbertConfig(**plbert_config)
+    bert = CustomAlbert(albert_base_configuration)
+    return bert

requirements.txt

@@ -0,0 +1,11 @@
+fugashi
+gradio
+mojimoji
+munch
+noisereduce
+phonemizer
+scipy
+torch
+transformers
+unicodedata2
+unidic-lite