Add LJSpeech model

app.py

@@ -1,5 +1,7 @@
 import gradio as gr
 import styletts2importable
+import ljspeechimportable
+import torch
 theme = gr.themes.Base(
     font=[gr.themes.GoogleFont('Libre Franklin'), gr.themes.GoogleFont('Public Sans'), 'system-ui', 'sans-serif'],
 )
@@ -20,7 +22,32 @@ def synthesize(text, voice):
         raise gr.Error("Text must be under 500 characters")
     v = voice.lower()
     return (24000, styletts2importable.inference(text, voices[v], alpha=0.3, beta=0.7, diffusion_steps=7, embedding_scale=1))
+def ljsynthesize(text):
+    if text.strip() == "":
+        raise gr.Error("You must enter some text")
+    if len(text) > 500:
+        raise gr.Error("Text must be under 500 characters")
+    noise = torch.randn(1,1,256).to('cuda' if torch.cuda.is_available() else 'cpu')
+    return (24000, ljspeechimportable.inference(text, noise, diffusion_steps=7, embedding_scale=1))
+
 
+with gr.Blocks() as vctk:
+    with gr.Row():
+        with gr.Column(scale=1):
+            inp = gr.Textbox(label="Text", info="What would you like StyleTTS 2 to read? It works better on full sentences.", interactive=True)
+            voice = gr.Dropdown(['Angie', 'Daniel', 'Tom', 'LJ', 'Pat', 'Tom', 'Dotrice', 'Mouse', 'William'], label="Voice", info="Select a voice. We use some voices from Tortoise TTS.", value='Tom', interactive=True)
+        with gr.Column(scale=1):
+            btn = gr.Button("Synthesize", variant="primary")
+            audio = gr.Audio(interactive=False, label="Synthesized Audio")
+            btn.click(synthesize, inputs=[inp, voice], outputs=[audio], concurrency_limit=4)
+with gr.Blocks() as lj:
+    with gr.Row():
+        with gr.Column(scale=1):
+            ljinp = gr.Textbox(label="Text", info="What would you like StyleTTS 2 to read? It works better on full sentences.", interactive=True)
+        with gr.Column(scale=1):
+            ljbtn = gr.Button("Synthesize", variant="primary")
+            ljaudio = gr.Audio(interactive=False, label="Synthesized Audio")
+            ljbtn.click(ljsynthesize, inputs=[ljinp], outputs=[ljaudio], concurrency_limit=4)
 with gr.Blocks(title="StyleTTS 2", css="footer{display:none !important}", theme=theme) as demo:
     gr.Markdown("""# StyleTTS 2
 
@@ -34,15 +61,7 @@ This space does NOT allow voice cloning. We use some default voice from Tortoise
 
 Is there a long queue on this space? Duplicate it and add a GPU to skip the wait!""")
     gr.DuplicateButton("Duplicate Space")
-    with gr.Row():
-        with gr.Column(scale=1):
-            inp = gr.Textbox(label="Text", info="What would you like StyleTTS 2 to read? It works better on full sentences.", interactive=True)
-            voice = gr.Dropdown(['Angie', 'Daniel', 'Tom', 'LJ', 'Pat', 'Tom', 'Dotrice', 'Mouse', 'William'], label="Voice", info="Select a voice. We use some voices from Tortoise TTS.", value='Tom', interactive=True)
-        with gr.Column(scale=1):
-            btn = gr.Button("Synthesize", variant="primary")
-            audio = gr.Audio(interactive=False, label="Synthesized Audio")
-            btn.click(synthesize, inputs=[inp, voice], outputs=[audio], concurrency_limit=4)
-    
+    gr.TabbedInterface([vctk, lj], ['Multi-Voice', 'LJSpeech'])
 if __name__ == "__main__":
     demo.queue(api_open=False, max_size=15).launch(show_api=False)
 

ljspeechimportable.py

@@ -0,0 +1,226 @@
+from cached_path import cached_path
+from dp.phonemizer import Phonemizer
+
+
+import torch
+torch.manual_seed(0)
+torch.backends.cudnn.benchmark = False
+torch.backends.cudnn.deterministic = True
+
+import random
+random.seed(0)
+
+import numpy as np
+np.random.seed(0)
+
+import nltk
+nltk.download('punkt')
+
+# load packages
+import time
+import random
+import yaml
+from munch import Munch
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torchaudio
+import librosa
+from nltk.tokenize import word_tokenize
+
+from models import *
+from utils import *
+from text_utils import TextCleaner
+textclenaer = TextCleaner()
+
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+to_mel = torchaudio.transforms.MelSpectrogram(
+    n_mels=80, n_fft=2048, win_length=1200, hop_length=300)
+mean, std = -4, 4
+
+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 preprocess(wave):
+    wave_tensor = torch.from_numpy(wave).float()
+    mel_tensor = to_mel(wave_tensor)
+    mel_tensor = (torch.log(1e-5 + mel_tensor.unsqueeze(0)) - mean) / std
+    return mel_tensor
+
+def compute_style(ref_dicts):
+    reference_embeddings = {}
+    for key, path in ref_dicts.items():
+        wave, sr = librosa.load(path, sr=24000)
+        audio, index = librosa.effects.trim(wave, top_db=30)
+        if sr != 24000:
+            audio = librosa.resample(audio, sr, 24000)
+        mel_tensor = preprocess(audio).to(device)
+
+        with torch.no_grad():
+            ref = model.style_encoder(mel_tensor.unsqueeze(1))
+        reference_embeddings[key] = (ref.squeeze(1), audio)
+
+    return reference_embeddings
+
+# load phonemizer
+# import phonemizer
+# global_phonemizer = phonemizer.backend.EspeakBackend(language='en-us', preserve_punctuation=True, with_stress=True, words_mismatch='ignore')
+
+phonemizer = Phonemizer.from_checkpoint(str(cached_path('https://public-asai-dl-models.s3.eu-central-1.amazonaws.com/DeepPhonemizer/en_us_cmudict_ipa_forward.pt')))
+
+
+config = yaml.safe_load(open(str(cached_path('hf://yl4579/StyleTTS2-LJSpeech/Models/LJSpeech/config.yml'))))
+
+# load pretrained ASR model
+ASR_config = config.get('ASR_config', False)
+ASR_path = config.get('ASR_path', False)
+text_aligner = load_ASR_models(ASR_path, ASR_config)
+
+# load pretrained F0 model
+F0_path = config.get('F0_path', False)
+pitch_extractor = load_F0_models(F0_path)
+
+# load BERT model
+from Utils.PLBERT.util import load_plbert
+BERT_path = config.get('PLBERT_dir', False)
+plbert = load_plbert(BERT_path)
+
+model = build_model(recursive_munch(config['model_params']), text_aligner, pitch_extractor, plbert)
+_ = [model[key].eval() for key in model]
+_ = [model[key].to(device) for key in model]
+
+# params_whole = torch.load("Models/LJSpeech/epoch_2nd_00100.pth", map_location='cpu')
+params_whole = torch.load(str(cached_path('hf://yl4579/StyleTTS2-LJSpeech/Models/LJSpeech/epoch_2nd_00100.pth')), map_location='cpu')
+params = params_whole['net']
+
+for key in model:
+    if key in params:
+        print('%s loaded' % key)
+        try:
+            model[key].load_state_dict(params[key])
+        except:
+            from collections import OrderedDict
+            state_dict = params[key]
+            new_state_dict = OrderedDict()
+            for k, v in state_dict.items():
+                name = k[7:] # remove `module.`
+                new_state_dict[name] = v
+            # load params
+            model[key].load_state_dict(new_state_dict, strict=False)
+#             except:
+#                 _load(params[key], model[key])
+_ = [model[key].eval() for key in model]
+
+from Modules.diffusion.sampler import DiffusionSampler, ADPM2Sampler, KarrasSchedule
+
+sampler = DiffusionSampler(
+    model.diffusion.diffusion,
+    sampler=ADPM2Sampler(),
+    sigma_schedule=KarrasSchedule(sigma_min=0.0001, sigma_max=3.0, rho=9.0), # empirical parameters
+    clamp=False
+)
+
+def inference(text, noise, diffusion_steps=5, embedding_scale=1):
+    text = text.strip()
+    text = text.replace('"', '')
+    ps = phonemizer([text], lang='en_us')
+    ps = word_tokenize(ps[0])
+    ps = ' '.join(ps)
+
+    tokens = textclenaer(ps)
+    tokens.insert(0, 0)
+    tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)
+
+    with torch.no_grad():
+        input_lengths = torch.LongTensor([tokens.shape[-1]]).to(tokens.device)
+        text_mask = length_to_mask(input_lengths).to(tokens.device)
+
+        t_en = model.text_encoder(tokens, input_lengths, text_mask)
+        bert_dur = model.bert(tokens, attention_mask=(~text_mask).int())
+        d_en = model.bert_encoder(bert_dur).transpose(-1, -2)
+
+        s_pred = sampler(noise,
+              embedding=bert_dur[0].unsqueeze(0), num_steps=diffusion_steps,
+              embedding_scale=embedding_scale).squeeze(0)
+
+        s = s_pred[:, 128:]
+        ref = s_pred[:, :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)
+        pred_dur = torch.round(duration.squeeze()).clamp(min=1)
+
+        pred_dur[-1] += 5
+
+        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)
+
+        # encode prosody
+        en = (d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0).to(device))
+        F0_pred, N_pred = model.predictor.F0Ntrain(en, s)
+        out = model.decoder((t_en @ pred_aln_trg.unsqueeze(0).to(device)),
+                                F0_pred, N_pred, ref.squeeze().unsqueeze(0))
+
+    return out.squeeze().cpu().numpy()
+
+def LFinference(text, s_prev, noise, alpha=0.7, diffusion_steps=5, embedding_scale=1):
+  text = text.strip()
+  text = text.replace('"', '')
+  ps = phonemizer([text], lang='en_us')
+  ps = word_tokenize(ps[0])
+  ps = ' '.join(ps)
+
+  tokens = textclenaer(ps)
+  tokens.insert(0, 0)
+  tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)
+
+  with torch.no_grad():
+      input_lengths = torch.LongTensor([tokens.shape[-1]]).to(tokens.device)
+      text_mask = length_to_mask(input_lengths).to(tokens.device)
+
+      t_en = model.text_encoder(tokens, input_lengths, text_mask)
+      bert_dur = model.bert(tokens, attention_mask=(~text_mask).int())
+      d_en = model.bert_encoder(bert_dur).transpose(-1, -2)
+
+      s_pred = sampler(noise,
+            embedding=bert_dur[0].unsqueeze(0), num_steps=diffusion_steps,
+            embedding_scale=embedding_scale).squeeze(0)
+
+      if s_prev is not None:
+          # convex combination of previous and current style
+          s_pred = alpha * s_prev + (1 - alpha) * s_pred
+
+      s = s_pred[:, 128:]
+      ref = s_pred[:, :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)
+      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)
+
+      # encode prosody
+      en = (d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0).to(device))
+      F0_pred, N_pred = model.predictor.F0Ntrain(en, s)
+      out = model.decoder((t_en @ pred_aln_trg.unsqueeze(0).to(device)),
+                              F0_pred, N_pred, ref.squeeze().unsqueeze(0))
+
+  return out.squeeze().cpu().numpy(), s_pred