Merge pull request #4 from PyThaiNLP/add-tts-thai

Add lunarlist_model

Author: wannaphong

notebook/use_lunarlist_model.ipynb

@@ -2,13 +2,13 @@
 """
 PyThaiTTS
 """
-__version__ = "0.1.1"
+__version__ = "0.2.0"
 
 
 class TTS:
-    def __init__(self, pretrained="khanomtan", mode="last_checkpoint", version="1.0") -> None:
+    def __init__(self, pretrained="khanomtan", mode="last_checkpoint", version="1.0", device:str="cpu") -> None:
         """
-        :param str pretrained: TTS pretrained (khanomtan)
+        :param str pretrained: TTS pretrained (khanomtan, lunarlist)
         :param str mode: pretrained mode
         :param str version: model version (default is 1.0 or 1.1)
 
@@ -18,9 +18,14 @@ def __init__(self, pretrained="khanomtan", mode="last_checkpoint", version="1.0"
         
         You can see more about khanomtan tts at `https://github.com/wannaphong/KhanomTan-TTS-v1.0 <https://github.com/wannaphong/KhanomTan-TTS-v1.0>`_
         and `https://github.com/wannaphong/KhanomTan-TTS-v1.1 <https://github.com/wannaphong/KhanomTan-TTS-v1.1>`_
+        
+        For lunarlist tts model, you must to install nemo before use the model by pip install nemo_toolkit['tts'].
+        You can see more about lunarlist tts at `https://link.medium.com/OpPjQis6wBb <https://link.medium.com/OpPjQis6wBb>`_
+        
         """
         self.pretrained = pretrained
         self.mode = mode
+        self.device = device
         self.load_pretrained(version=version)
 
     def load_pretrained(self,version):
@@ -30,6 +35,9 @@ def load_pretrained(self,version):
         if self.pretrained == "khanomtan":
             from pythaitts.pretrained import KhanomTan
             self.model = KhanomTan(mode=self.mode, version=version)
+        elif self.pretrained == "lunarlist":
+            from pythaitts.pretrained import LunarlistModel
+            self.model = LunarlistModel(mode=self.mode, device=self.device)
         else:
             raise NotImplemented(
                 "PyThaiTTS doesn't support %s pretrained." % self.pretrained
@@ -45,6 +53,8 @@ def tts(self, text: str, speaker_idx: str = "Linda", language_idx: str = "th-th"
         :param str return_type: return type (default is file)
         :param str filename: path filename for save wav file if return_type is file.
         """
+        if self.pretrained == "lunarlist":
+            return self.model(text=text,return_type=return_type,filename=filename)
         return self.model(
             text=text,
             speaker_idx=speaker_idx,

pythaitts/__init__.py

@@ -1,6 +1,8 @@
 # -*- coding: utf-8 -*-
 from pythaitts.pretrained.khanomtan_tts import KhanomTan
+from pythaitts.pretrained.lunarlist_model import LunarlistModel
 
 __all__ = [
-    "KhanomTan"
+    "KhanomTan",
+    "LunarlistModel"
 ]

pythaitts/pretrained/__init__.py

@@ -0,0 +1,44 @@
+# -*- coding: utf-8 -*-
+"""
+Lunarlist TTS model
+
+You can see more about lunarlist tts at `https://link.medium.com/OpPjQis6wBb <https://link.medium.com/OpPjQis6wBb>`_
+"""
+import tempfile
+import torch
+
+
+class LunarlistModel:
+    def __init__(self, mode:str="last_checkpoint", device:str="cpu") -> None:
+        try:
+            from nemo.collections.tts.models import UnivNetModel
+        except ImportError:
+            raise ImportError("You must to install nemo by pip install nemo_toolkit['tts'] before use this model.")
+        self.mode = mode
+        self.device = device
+        self.vcoder_model = UnivNetModel.from_pretrained(model_name="tts_en_libritts_univnet").to(self.device)
+        self.load_synthesizer(self.mode)
+    def load_synthesizer(self, mode:str):
+        from nemo.collections.tts.models import Tacotron2Model
+        if mode=="last_checkpoint":
+            self.model = Tacotron2Model.from_pretrained("lunarlist/tts-thai-last-step").to(self.device)
+        else:
+            self.model = Tacotron2Model.from_pretrained("lunarlist/tts-thai").to(self.device)
+        self.dict_idx={k:i for i,k in enumerate(self.model.hparams["cfg"]['labels'])}
+    def tts(self, text: str):
+        parsed2=torch.Tensor([[66]+[self.dict_idx[i] for i in text if i]+[67]]).int().to(self.device)
+        spectrogram2 = self.model.generate_spectrogram(tokens=parsed2)
+        audio2 = self.vcoder_model.convert_spectrogram_to_audio(spec=spectrogram2)
+        return audio2.to('cpu').detach().numpy()
+    def __call__(self, text: str,return_type: str = "file", filename: str = None):
+        wavs = self.tts(text)
+        if return_type == "waveform":
+            return wavs
+        import soundfile as sf
+        if filename != None:
+            sf.write(filename, wavs[0], 22050)
+            return filename
+        else:
+            with tempfile.NamedTemporaryFile(suffix = ".wav", delete = False) as fp:
+                fp.write(wavs[0])
+            return fp.name

pythaitts/pretrained/lunarlist_model.py

@@ -1,3 +1,4 @@
 TTS>=0.8.0
 pythainlp>=3.0.0
-huggingface_hub
+huggingface_hub
+torch

requirements.txt

@@ -9,7 +9,7 @@
 
 setup(
     name="PyThaiTTS",
-    version="0.1.1",
+    version="0.2.0",
     description="Open Source Thai Text-to-speech library in Python",
     long_description=readme,
     long_description_content_type="text/markdown",