diff --git a/README.md b/README.md index 56a0237..f0bd02a 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,11 @@ from pythaitts import TTS tts = TTS(pretrained="fastthaig2p") file = tts.tts("สวัสดีครับ", speaker_idx="thai_som", filename="output.wav") +# Use MakHamTTS (Paradee Thai INT8) +# Sample Rate is 24000 Hz +tts = TTS(pretrained="makhamtts") +file = tts.tts("สวัสดีค่ะ ยินดีต้อนรับสู่บริการของเรา", speaker_idx="paradee_thai", filename="output.wav") + # Use Lunarlist ONNX # Sample Rate is 22050 Hz tts = TTS(pretrained="lunarlist_onnx") diff --git a/pythaitts/__init__.py b/pythaitts/__init__.py index 5710e7f..994a928 100644 --- a/pythaitts/__init__.py +++ b/pythaitts/__init__.py @@ -26,7 +26,7 @@ class TTS: def __init__(self, pretrained="fastthaig2p", mode="last_checkpoint", version="1.0", device:str="cpu", **kwargs) -> None: """ - :param str pretrained: TTS pretrained (fastthaig2p, lunarlist_onnx, khanomtan, lunarlist, vachana) + :param str pretrained: TTS pretrained (fastthaig2p, makhamtts, lunarlist_onnx, khanomtan, lunarlist, vachana) :param str mode: pretrained mode (fastthaig2p, lunarlist_onnx, and vachana don't support) :param str version: model version (default is 1.0 or 1.1) :param str device: device for running model. (fastthaig2p, lunarlist_onnx, and vachana support CPU only.) @@ -72,6 +72,9 @@ def load_pretrained(self, version="1.0", **kwargs): elif self.pretrained == "vachana": from pythaitts.pretrained.vachana_tts import VachanaTTS self.model = VachanaTTS() + elif self.pretrained in ("makhamtts", "MakHamTTS"): + from pythaitts.pretrained.makham_tts import MakHamTTS + self.model = MakHamTTS() elif self.pretrained in ("fastthaig2p", "FastThaiG2P"): from pythaitts.pretrained.fastthaig2p import FastThaiG2P self.model = FastThaiG2P(device=self.device, **kwargs) @@ -103,6 +106,10 @@ def tts(self, text: str, speaker_idx: str = "thai_som", language_idx: str = "th- if speaker_idx in ("thai_som", "Linda", None): speaker_idx = "th_f_1" return self.model(text=text,speaker_idx=speaker_idx,return_type=return_type,filename=filename, **kwargs) + elif self.pretrained in ("makhamtts", "MakHamTTS"): + if speaker_idx in ("thai_som", "Linda", None): + speaker_idx = "paradee_thai" + return self.model(text=text,speaker_idx=speaker_idx,return_type=return_type,filename=filename, **kwargs) elif self.pretrained in ("fastthaig2p", "FastThaiG2P"): if speaker_idx in ("Linda", None): speaker_idx = "thai_som" diff --git a/pythaitts/pretrained/makham_tts.py b/pythaitts/pretrained/makham_tts.py new file mode 100644 index 0000000..ad6f093 --- /dev/null +++ b/pythaitts/pretrained/makham_tts.py @@ -0,0 +1,78 @@ +# -*- coding: utf-8 -*- +import json +import tempfile +from pathlib import Path + +import numpy as np +import onnxruntime as ort +import soundfile as sf + +from pythaitts.pretrained.fastthaig2p import G2P, ipa_to_kokoro + + +_ASSET_DIR = Path(__file__).resolve().parent / "makhamtts" +_DEFAULT_MODEL_PATH = _ASSET_DIR / "paradee_thai_int8.onnx" +_DEFAULT_CONFIG_PATH = _ASSET_DIR / "config.json" + + +class MakHamTTS: + SUPPORTED_VOICES = ["paradee_thai"] + + def __init__(self, model_path=None, config_path=None) -> None: + self.model_path = Path(model_path) if model_path else _DEFAULT_MODEL_PATH + self.config_path = Path(config_path) if config_path else _DEFAULT_CONFIG_PATH + + if not self.model_path.exists(): + raise FileNotFoundError(f"MakHamTTS model file not found: {self.model_path}") + if not self.config_path.exists(): + raise FileNotFoundError(f"MakHamTTS config file not found: {self.config_path}") + + with open(self.config_path, "r", encoding="utf-8") as f: + cfg = json.load(f) + + self.vocab = cfg["vocab"] + self.sample_rate = cfg.get("sample_rate", 24000) + self.g2p = G2P() + self.session = ort.InferenceSession( + str(self.model_path), providers=["CPUExecutionProvider"] + ) + + def tts(self, text: str, speed: float = 1.0): + phonemes = ipa_to_kokoro(self.g2p.convert(text)) + ids = [0] + [self.vocab[c] for c in phonemes if c in self.vocab] + [0] + if len(ids) <= 2: + return np.array([], dtype=np.float32) + audio = self.session.run( + None, + { + "input_ids": np.array([ids], dtype=np.int64), + "speed": np.array([float(speed)], dtype=np.float32), + }, + )[0][0] + return audio + + def __call__( + self, + text: str, + speaker_idx: str = "paradee_thai", + return_type: str = "file", + filename: str = None, + **kwargs, + ): + if speaker_idx in ("thai_som", "Linda", None): + speaker_idx = "paradee_thai" + if speaker_idx not in self.SUPPORTED_VOICES: + raise ValueError( + f"Unsupported voice '{speaker_idx}'. Supported voices are: {', '.join(self.SUPPORTED_VOICES)}" + ) + + wav = self.tts(text, speed=kwargs.get("speed", 1.0)) + + if return_type == "waveform": + return wav + if filename is not None: + sf.write(filename, wav, self.sample_rate) + return filename + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as fp: + sf.write(fp.name, wav, self.sample_rate) + return fp.name diff --git a/pythaitts/pretrained/makhamtts/config.json b/pythaitts/pretrained/makhamtts/config.json new file mode 100644 index 0000000..bfe2792 --- /dev/null +++ b/pythaitts/pretrained/makhamtts/config.json @@ -0,0 +1,179 @@ +{ + "istftnet": { + "upsample_kernel_sizes": [ + 20, + 12 + ], + "upsample_rates": [ + 10, + 6 + ], + "gen_istft_hop_size": 5, + "gen_istft_n_fft": 20, + "resblock_dilation_sizes": [ + [ + 1, + 3, + 5 + ], + [ + 1, + 3, + 5 + ], + [ + 1, + 3, + 5 + ] + ], + "resblock_kernel_sizes": [ + 3, + 7, + 11 + ], + "upsample_initial_channel": 512 + }, + "dim_in": 64, + "dropout": 0.2, + "hidden_dim": 512, + "max_conv_dim": 512, + "max_dur": 50, + "multispeaker": true, + "n_layer": 3, + "n_mels": 80, + "n_token": 178, + "style_dim": 128, + "text_encoder_kernel_size": 5, + "plbert": { + "hidden_size": 768, + "num_attention_heads": 12, + "intermediate_size": 2048, + "max_position_embeddings": 512, + "num_hidden_layers": 12, + "dropout": 0.1 + }, + "vocab": { + ";": 1, + ":": 2, + ",": 3, + ".": 4, + "!": 5, + "?": 6, + "—": 9, + "…": 10, + "\"": 11, + "(": 12, + ")": 13, + "“": 14, + "”": 15, + " ": 16, + "̃": 17, + "ʣ": 18, + "ʥ": 19, + "ʦ": 20, + "ʨ": 21, + "ᵝ": 22, + "ꭧ": 23, + "A": 24, + "I": 25, + "O": 31, + "Q": 33, + "S": 35, + "T": 36, + "W": 39, + "Y": 41, + "ᵊ": 42, + "a": 43, + "b": 44, + "c": 45, + "d": 46, + "e": 47, + "f": 48, + "h": 50, + "i": 51, + "j": 52, + "k": 53, + "l": 54, + "m": 55, + "n": 56, + "o": 57, + "p": 58, + "q": 59, + "r": 60, + "s": 61, + "t": 62, + "u": 63, + "v": 64, + "w": 65, + "x": 66, + "y": 67, + "z": 68, + "ɑ": 69, + "ɐ": 70, + "ɒ": 71, + "æ": 72, + "β": 75, + "ɔ": 76, + "ɕ": 77, + "ç": 78, + "ɖ": 80, + "ð": 81, + "ʤ": 82, + "ə": 83, + "ɚ": 85, + "ɛ": 86, + "ɜ": 87, + "ɟ": 90, + "ɡ": 92, + "ɥ": 99, + "ɨ": 101, + "ɪ": 102, + "ʝ": 103, + "ɯ": 110, + "ɰ": 111, + "ŋ": 112, + "ɳ": 113, + "ɲ": 114, + "ɴ": 115, + "ø": 116, + "ɸ": 118, + "θ": 119, + "œ": 120, + "ɹ": 123, + "ɾ": 125, + "ɻ": 126, + "ʁ": 128, + "ɽ": 129, + "ʂ": 130, + "ʃ": 131, + "ʈ": 132, + "ʧ": 133, + "ʊ": 135, + "ʋ": 136, + "ʌ": 138, + "ɣ": 139, + "ɤ": 140, + "χ": 142, + "ʎ": 143, + "ʒ": 147, + "ʔ": 148, + "ˈ": 156, + "ˌ": 157, + "ː": 158, + "ʰ": 162, + "ʲ": 164, + "↓": 169, + "→": 171, + "↗": 172, + "↘": 173, + "ᵻ": 177, + "↑": 170 + }, + "model_type": "paradee", + "name": "Paradee-Thai-Female", + "language": "th", + "sample_rate": 24000, + "teacher": "kokoro_thai", + "teacher_voice": "newvoice_female" +} \ No newline at end of file diff --git a/pythaitts/pretrained/makhamtts/paradee_thai_int8.onnx b/pythaitts/pretrained/makhamtts/paradee_thai_int8.onnx new file mode 100644 index 0000000..34e69f1 Binary files /dev/null and b/pythaitts/pretrained/makhamtts/paradee_thai_int8.onnx differ diff --git a/setup.py b/setup.py index 8cd6564..1d53b75 100644 --- a/setup.py +++ b/setup.py @@ -19,6 +19,7 @@ packages=find_packages(), package_data={ "pythaitts.pretrained.fastthaig2p": ["data/*", "data/fallback/*"], + "pythaitts.pretrained": ["makhamtts/*.onnx", "makhamtts/*.json"], }, include_package_data=True, python_requires=">=3.6", diff --git a/tests/test_makhamtts.py b/tests/test_makhamtts.py new file mode 100644 index 0000000..c219865 --- /dev/null +++ b/tests/test_makhamtts.py @@ -0,0 +1,57 @@ +# -*- coding: utf-8 -*- +import unittest +from unittest.mock import Mock, patch + +import numpy as np + +from pythaitts import TTS + + +class TestMakHamTTSIntegration(unittest.TestCase): + @patch("pythaitts.pretrained.makham_tts.MakHamTTS") + def test_makhamtts_model_initialization(self, mock_makham_cls): + mock_makham_cls.return_value = Mock() + tts = TTS(pretrained="makhamtts") + self.assertIsNotNone(tts.model) + self.assertEqual(tts.pretrained, "makhamtts") + + @patch("pythaitts.pretrained.makham_tts.MakHamTTS") + def test_makhamtts_tts_call(self, mock_makham_cls): + mock_instance = Mock() + mock_instance.return_value = "/tmp/output.wav" + mock_makham_cls.return_value = mock_instance + + tts = TTS(pretrained="makhamtts") + result = tts.tts("สวัสดีค่ะ", speaker_idx="paradee_thai", filename="/tmp/test.wav") + + self.assertEqual(result, "/tmp/output.wav") + call_args = mock_instance.call_args + self.assertEqual(call_args.kwargs["text"], "สวัสดีค่ะ") + self.assertEqual(call_args.kwargs["speaker_idx"], "paradee_thai") + self.assertEqual(call_args.kwargs["filename"], "/tmp/test.wav") + self.assertEqual(call_args.kwargs["return_type"], "file") + + @patch("pythaitts.pretrained.makham_tts.ipa_to_kokoro") + @patch("pythaitts.pretrained.makham_tts.G2P") + @patch("pythaitts.pretrained.makham_tts.ort.InferenceSession") + def test_makhamtts_waveform_return( + self, mock_session_cls, mock_g2p_cls, mock_ipa_to_kokoro + ): + from pythaitts.pretrained.makham_tts import MakHamTTS + + mock_session = Mock() + mock_session.run.return_value = [np.array([[0.1, 0.2]], dtype=np.float32)] + mock_session_cls.return_value = mock_session + mock_g2p = Mock() + mock_g2p.convert.return_value = "/dummy/" + mock_g2p_cls.return_value = mock_g2p + mock_ipa_to_kokoro.return_value = "abc" + + model = MakHamTTS() + wav = model("สวัสดีค่ะ", return_type="waveform") + + self.assertTrue(np.array_equal(wav, np.array([0.1, 0.2], dtype=np.float32))) + + +if __name__ == "__main__": + unittest.main()