diff --git a/onnx_infer.py b/onnx_infer.py new file mode 100644 index 0000000..81df5bc --- /dev/null +++ b/onnx_infer.py @@ -0,0 +1,71 @@ +from onnx_modules.V220_OnnxInference import OnnxInferenceSession +import numpy as np +Session = OnnxInferenceSession( + { + "enc" : "onnx/BertVits2.2PT/BertVits2.2PT_enc_p.onnx", + "emb_g" : "onnx/BertVits2.2PT/BertVits2.2PT_emb.onnx", + "dp" : "onnx/BertVits2.2PT/BertVits2.2PT_dp.onnx", + "sdp" : "onnx/BertVits2.2PT/BertVits2.2PT_sdp.onnx", + "flow" : "onnx/BertVits2.2PT/BertVits2.2PT_flow.onnx", + "dec" : "onnx/BertVits2.2PT/BertVits2.2PT_dec.onnx" + }, + Providers = ["CPUExecutionProvider"] + ) + +#这里的输入和原版是一样的,只需要在原版预处理结果出来之后加上.numpy()即可 +x = np.expand_dims( + np.array( + [ + 0, + 97, + 0, + 8, + 0, + 78, + 0, + 8, + 0, + 76, + 0, + 37, + 0, + 40, + 0, + 97, + 0, + 8, + 0, + 23, + 0, + 8, + 0, + 74, + 0, + 26, + 0, + 104, + 0, + ] + ), + 0 +) +tone = np.zeros_like(x) +language = np.zeros_like(x) +sid = np.array([0]) +bert = np.random.randn(x.shape[1], 1024) +ja_bert = np.random.randn(x.shape[1], 1024) +en_bert = np.random.randn(x.shape[1], 1024) +emo = np.random.randn(512, 1) + +audio = Session( + x, + tone, + language, + bert, + ja_bert, + en_bert, + emo, + sid +) + +print(audio) \ No newline at end of file diff --git a/onnx_modules/V220_OnnxInference/__init__.py b/onnx_modules/V220_OnnxInference/__init__.py new file mode 100644 index 0000000..74256d2 --- /dev/null +++ b/onnx_modules/V220_OnnxInference/__init__.py @@ -0,0 +1,92 @@ +import numpy as np +import onnxruntime as ort + +def convert_pad_shape(pad_shape): + layer = pad_shape[::-1] + pad_shape = [item for sublist in layer for item in sublist] + return pad_shape + +def sequence_mask(length, max_length=None): + if max_length is None: + max_length = length.max() + x = np.arange(max_length, dtype=length.dtype) + return np.expand_dims(x, 0) < np.expand_dims(length, 1) + +def generate_path(duration, mask): + """ + duration: [b, 1, t_x] + mask: [b, 1, t_y, t_x] + """ + + b, _, t_y, t_x = mask.shape + cum_duration = np.cumsum(duration, -1) + + cum_duration_flat = cum_duration.reshape(b * t_x) + path = sequence_mask(cum_duration_flat, t_y) + path = path.reshape(b, t_x, t_y) + path = path ^ np.pad(path, ((0, 0), (1, 0), (0, 0)))[:, :-1] + path = np.expand_dims(path, 1).transpose(0, 1, 3, 2) + return path + +class OnnxInferenceSession(): + def __init__(self, path, Providers = ["CPUExecutionProvider"]): + self.enc = ort.InferenceSession(path["enc"], providers=Providers) + self.emb_g = ort.InferenceSession(path["emb_g"], providers=Providers) + self.dp = ort.InferenceSession(path["dp"], providers=Providers) + self.sdp = ort.InferenceSession(path["sdp"], providers=Providers) + self.flow = ort.InferenceSession(path["flow"], providers=Providers) + self.dec = ort.InferenceSession(path["dec"], providers=Providers) + + def __call__( + self, + seq, + tone, + language, + bert_zh, + bert_jp, + bert_en, + emo, + sid, + seed = 114514, + seq_noise_scale = 0.8, + sdp_noise_scale = 0.6, + length_scale = 1., + sdp_ratio = 0. + ): + g = self.emb_g.run(None, {'sid': sid.astype(np.int64),})[0] + g = np.expand_dims(g, -1) + enc_rtn = self.enc.run( + None, + { + "x" : seq.astype(np.int64), + "t" : tone.astype(np.int64), + "language" : language.astype(np.int64), + "bert_0" : bert_zh.astype(np.float32), + "bert_1" : bert_jp.astype(np.float32), + "bert_2" : bert_en.astype(np.float32), + "emo" : emo.astype(np.float32), + "g" : g.astype(np.float32) + }) + x, m_p, logs_p, x_mask = enc_rtn[0], enc_rtn[1], enc_rtn[2], enc_rtn[3] + np.random.seed(seed) + zinput = np.random.randn(x.shape[0], 2, x.shape[2]) * sdp_noise_scale + logw = self.sdp.run(None, {"x" : x, "x_mask" : x_mask, "zin" : zinput.astype(np.float32), "g" : g})[0] * (sdp_ratio) + \ + self.dp.run(None, {"x" : x, "x_mask" : x_mask, "g" : g})[0] * (1 - sdp_ratio) + w = np.exp(logw) * x_mask * length_scale + w_ceil = np.ceil(w) + y_lengths = np.clip(np.sum(w_ceil, (1, 2)), a_min=1., a_max=100000).astype(np.int64) + y_mask = np.expand_dims(sequence_mask(y_lengths, None), 1) + attn_mask = np.expand_dims(x_mask, 2) * np.expand_dims(y_mask, -1) + attn = generate_path(w_ceil, attn_mask) + m_p = np.matmul(attn.squeeze(1), m_p.transpose(0, 2, 1)).transpose( + 0, 2, 1 + ) # [b, t', t], [b, t, d] -> [b, d, t'] + logs_p = np.matmul(attn.squeeze(1), logs_p.transpose(0, 2, 1)).transpose( + 0, 2, 1 + ) # [b, t', t], [b, t, d] -> [b, d, t'] + + z_p = m_p + np.random.randn(m_p.shape[0], m_p.shape[1], m_p.shape[2]) * np.exp(logs_p) * seq_noise_scale + + z = self.flow.run(None, {"z_p" : z_p.astype(np.float32), "y_mask" : y_mask.astype(np.float32), "g": g})[0] + + return self.dec.run(None, {"z_in" : z.astype(np.float32), "g": g})[0] \ No newline at end of file