Export MeloTTS to ONNX (#1129)
This commit is contained in:
Fangjun Kuang
256
scripts/melo-tts/export-onnx.py
Executable file
256
scripts/melo-tts/export-onnx.py
Executable file
@@ -0,0 +1,256 @@
|
||||
#!/usr/bin/env python3
|
||||
from typing import Any, Dict
|
||||
|
||||
import onnx
|
||||
import torch
|
||||
from melo.api import TTS
|
||||
from melo.text import language_id_map, language_tone_start_map
|
||||
from melo.text.chinese import pinyin_to_symbol_map
|
||||
from pypinyin import Style, lazy_pinyin, phrases_dict, pinyin_dict
|
||||
|
||||
for k, v in pinyin_to_symbol_map.items():
|
||||
pinyin_to_symbol_map[k] = v.split()
|
||||
|
||||
|
||||
def get_initial_final_tone(word: str):
|
||||
initials = lazy_pinyin(word, neutral_tone_with_five=True, style=Style.INITIALS)
|
||||
finals = lazy_pinyin(word, neutral_tone_with_five=True, style=Style.FINALS_TONE3)
|
||||
|
||||
ans_phone = []
|
||||
ans_tone = []
|
||||
|
||||
for c, v in zip(initials, finals):
|
||||
raw_pinyin = c + v
|
||||
v_without_tone = v[:-1]
|
||||
try:
|
||||
tone = v[-1]
|
||||
except:
|
||||
print("skip", word, initials, finals)
|
||||
return [], []
|
||||
|
||||
pinyin = c + v_without_tone
|
||||
assert tone in "12345"
|
||||
|
||||
if c:
|
||||
v_rep_map = {
|
||||
"uei": "ui",
|
||||
"iou": "iu",
|
||||
"uen": "un",
|
||||
}
|
||||
if v_without_tone in v_rep_map.keys():
|
||||
pinyin = c + v_rep_map[v_without_tone]
|
||||
else:
|
||||
pinyin_rep_map = {
|
||||
"ing": "ying",
|
||||
"i": "yi",
|
||||
"in": "yin",
|
||||
"u": "wu",
|
||||
}
|
||||
if pinyin in pinyin_rep_map.keys():
|
||||
pinyin = pinyin_rep_map[pinyin]
|
||||
else:
|
||||
single_rep_map = {
|
||||
"v": "yu",
|
||||
"e": "e",
|
||||
"i": "y",
|
||||
"u": "w",
|
||||
}
|
||||
if pinyin[0] in single_rep_map.keys():
|
||||
pinyin = single_rep_map[pinyin[0]] + pinyin[1:]
|
||||
# print(word, initials, finals, pinyin)
|
||||
|
||||
if pinyin not in pinyin_to_symbol_map:
|
||||
print("skip", pinyin, word, c, v, raw_pinyin)
|
||||
continue
|
||||
phone = pinyin_to_symbol_map[pinyin]
|
||||
ans_phone += phone
|
||||
ans_tone += [tone] * len(phone)
|
||||
|
||||
return ans_phone, ans_tone
|
||||
|
||||
|
||||
def generate_tokens(symbol_list):
|
||||
with open("tokens.txt", "w", encoding="utf-8") as f:
|
||||
for i, s in enumerate(symbol_list):
|
||||
f.write(f"{s} {i}\n")
|
||||
|
||||
|
||||
def generate_lexicon():
|
||||
word_dict = pinyin_dict.pinyin_dict
|
||||
phrases = phrases_dict.phrases_dict
|
||||
with open("lexicon.txt", "w", encoding="utf-8") as f:
|
||||
for key in word_dict:
|
||||
if not (0x4E00 <= key <= 0x9FA5):
|
||||
continue
|
||||
w = chr(key)
|
||||
phone, tone = get_initial_final_tone(w)
|
||||
if not phone:
|
||||
continue
|
||||
phone = " ".join(phone)
|
||||
tone = " ".join(tone)
|
||||
f.write(f"{w} {phone} {tone}\n")
|
||||
|
||||
for w in phrases:
|
||||
phone, tone = get_initial_final_tone(w)
|
||||
if not phone:
|
||||
continue
|
||||
assert len(phone) == len(tone), (len(phone), len(tone), phone, tone)
|
||||
phone = " ".join(phone)
|
||||
tone = " ".join(tone)
|
||||
f.write(f"{w} {phone} {tone}\n")
|
||||
|
||||
|
||||
def add_meta_data(filename: str, meta_data: Dict[str, Any]):
|
||||
"""Add meta data to an ONNX model. It is changed in-place.
|
||||
|
||||
Args:
|
||||
filename:
|
||||
Filename of the ONNX model to be changed.
|
||||
meta_data:
|
||||
Key-value pairs.
|
||||
"""
|
||||
model = onnx.load(filename)
|
||||
while len(model.metadata_props):
|
||||
model.metadata_props.pop()
|
||||
|
||||
for key, value in meta_data.items():
|
||||
meta = model.metadata_props.add()
|
||||
meta.key = key
|
||||
meta.value = str(value)
|
||||
|
||||
onnx.save(model, filename)
|
||||
|
||||
|
||||
class ModelWrapper(torch.nn.Module):
|
||||
def __init__(self, model: "SynthesizerTrn"):
|
||||
super().__init__()
|
||||
self.model = model
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
x_lengths,
|
||||
tones,
|
||||
lang_id,
|
||||
bert,
|
||||
ja_bert,
|
||||
sid,
|
||||
noise_scale,
|
||||
length_scale,
|
||||
noise_scale_w,
|
||||
max_len=None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
x: A 1-D array of dtype np.int64. Its shape is (token_numbers,)
|
||||
tones: A 1-D array of dtype np.int64. Its shape is (token_numbers,)
|
||||
lang_id: A 1-D array of dtype np.int64. Its shape is (token_numbers,)
|
||||
sid: an integer
|
||||
"""
|
||||
return self.model.infer(
|
||||
x=x,
|
||||
x_lengths=x_lengths,
|
||||
sid=sid,
|
||||
tone=tones,
|
||||
language=lang_id,
|
||||
bert=bert,
|
||||
ja_bert=ja_bert,
|
||||
noise_scale=noise_scale,
|
||||
noise_scale_w=noise_scale_w,
|
||||
length_scale=length_scale,
|
||||
)[0]
|
||||
|
||||
|
||||
def main():
|
||||
generate_lexicon()
|
||||
|
||||
language = "ZH"
|
||||
model = TTS(language=language, device="cpu")
|
||||
|
||||
generate_tokens(model.hps["symbols"])
|
||||
|
||||
torch_model = ModelWrapper(model.model)
|
||||
|
||||
opset_version = 13
|
||||
x = torch.randint(low=0, high=10, size=(60,), dtype=torch.int64)
|
||||
print(x.shape)
|
||||
x_lengths = torch.tensor([x.size(0)], dtype=torch.int64)
|
||||
sid = torch.tensor([1], dtype=torch.int64)
|
||||
tones = torch.zeros_like(x)
|
||||
lang_id = torch.ones_like(x)
|
||||
noise_scale = torch.tensor([1.0], dtype=torch.float32)
|
||||
length_scale = torch.tensor([1.0], dtype=torch.float32)
|
||||
noise_scale_w = torch.tensor([1.0], dtype=torch.float32)
|
||||
|
||||
bert = torch.zeros(1024, x.shape[0], dtype=torch.float32)
|
||||
ja_bert = torch.zeros(768, x.shape[0], dtype=torch.float32)
|
||||
|
||||
x = x.unsqueeze(0)
|
||||
tones = tones.unsqueeze(0)
|
||||
lang_id = lang_id.unsqueeze(0)
|
||||
bert = bert.unsqueeze(0)
|
||||
ja_bert = ja_bert.unsqueeze(0)
|
||||
|
||||
filename = "model.onnx"
|
||||
|
||||
torch.onnx.export(
|
||||
torch_model,
|
||||
(
|
||||
x,
|
||||
x_lengths,
|
||||
tones,
|
||||
lang_id,
|
||||
bert,
|
||||
ja_bert,
|
||||
sid,
|
||||
noise_scale,
|
||||
length_scale,
|
||||
noise_scale_w,
|
||||
),
|
||||
filename,
|
||||
opset_version=opset_version,
|
||||
input_names=[
|
||||
"x",
|
||||
"x_lengths",
|
||||
"tones",
|
||||
"lang_id",
|
||||
"bert",
|
||||
"ja_bert",
|
||||
"sid",
|
||||
"noise_scale",
|
||||
"length_scale",
|
||||
"noise_scale_w",
|
||||
],
|
||||
output_names=["y"],
|
||||
dynamic_axes={
|
||||
"x": {0: "N", 1: "L"},
|
||||
"x_lengths": {0: "N"},
|
||||
"tones": {0: "N", 1: "L"},
|
||||
"lang_id": {0: "N", 1: "L"},
|
||||
"bert": {0: "N", 2: "L"},
|
||||
"ja_bert": {0: "N", 2: "L"},
|
||||
"y": {0: "N", 1: "S", 2: "T"},
|
||||
},
|
||||
)
|
||||
|
||||
meta_data = {
|
||||
"model_type": "melo-vits",
|
||||
"comment": "melo",
|
||||
"language": "Chinese + English",
|
||||
"add_blank": int(model.hps.data.add_blank),
|
||||
"n_speakers": 1,
|
||||
"sample_rate": model.hps.data.sampling_rate,
|
||||
"bert_dim": 1024,
|
||||
"ja_bert_dim": 768,
|
||||
"speaker_id": list(model.hps.data.spk2id.values())[0],
|
||||
"lang_id": language_id_map[model.language],
|
||||
"tone_start": language_tone_start_map[model.language],
|
||||
"url": "https://github.com/myshell-ai/MeloTTS",
|
||||
"license": "MIT license",
|
||||
"description": "MeloTTS is a high-quality multi-lingual text-to-speech library by MyShell.ai",
|
||||
}
|
||||
add_meta_data(filename, meta_data)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user