forked from FENRlR/MB-iSTFT-VITS2
-
Notifications
You must be signed in to change notification settings - Fork 0
/
onnx_export.py
104 lines (82 loc) · 2.83 KB
/
onnx_export.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
### from @nshmyrev's fork : https://github.com/alphacep/MB-iSTFT-VITS2/blob/main/export.py
import librosa
import matplotlib.pyplot as plt
import os
import json
import math
import requests
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
import commons
import utils
from data_utils import TextAudioLoader, TextAudioCollate, TextAudioSpeakerLoader, TextAudioSpeakerCollate
from models import SynthesizerTrn
from text.symbols import symbols
import numpy as np
from scipy.io.wavfile import write
import re
from scipy import signal
#- Variable section
PATH_TO_CONFIG = "./configs/PATH_TO_CONFIG.json" # path to config
PATH_TO_MODEL = "PATH_TO_G_XXXX.pth" # path to model
SPEAKER_ID = None # sid
SCALE_CONFIG = torch.FloatTensor([0.667, 1.0, 0.8]) # scales -> noise, noise_w, length
OPSET_VERSION = 15
hps = utils.get_hparams_from_file(PATH_TO_CONFIG)
if "use_mel_posterior_encoder" in hps.model.keys() and hps.model.use_mel_posterior_encoder == True:
print("Using mel posterior encoder for VITS2")
posterior_channels = 80 # vits2
hps.data.use_mel_posterior_encoder = True
else:
print("Using lin posterior encoder for VITS1")
posterior_channels = hps.data.filter_length // 2 + 1
hps.data.use_mel_posterior_encoder = False
net_g = SynthesizerTrn(
len(symbols),
posterior_channels,
hps.train.segment_size // hps.data.hop_length,
n_speakers=hps.data.n_speakers,
is_onnx=True, # !
**hps.model)
_ = utils.load_checkpoint(PATH_TO_MODEL, net_g, None)
num_symbols = net_g.n_vocab
num_speakers = net_g.n_speakers
def infer_forward(text, text_lengths, scales, sid=None):
noise_scale = scales[0]
length_scale = scales[1]
noise_scale_w = scales[2]
audio = net_g.infer(
text,
text_lengths,
noise_scale=noise_scale,
length_scale=length_scale,
noise_scale_w=noise_scale_w,
sid=sid,
)[0].unsqueeze(1)
return audio
with torch.no_grad():
net_g.dec.remove_weight_norm()
net_g.flow.remove_weight_norm() # Remove weightnorm in flows - a8d9f74
net_g.forward = infer_forward
net_g.eval()
# dummy initialization
dmy_text = torch.randint(low=0, high=num_symbols, size=(1, 50), dtype=torch.long)
dmy_text_length = torch.LongTensor([dmy_text.size(1)])
dummy_input = (dmy_text, dmy_text_length, SCALE_CONFIG, SPEAKER_ID) # infer_forward()
# Export
torch.onnx.export(
model=net_g,
args=dummy_input,
f="model.onnx",
verbose=True,
opset_version=OPSET_VERSION,
input_names=["input", "input_lengths", "scales", "sid"],
output_names=["output"],
dynamic_axes={
"input": {0: "batch_size", 1: "phonemes"},
"input_lengths": {0: "batch_size"},
"output": {0: "batch_size", 1: "time"},
},
)