mirror of
https://github.com/storytold/vits-finetuning.git
synced 2026-10-09 00:09:52 +00:00
ZDisket
137 lines
3.8 KiB
Python
137 lines
3.8 KiB
Python
import os
|
|
import json
|
|
import argparse
|
|
|
|
import torch
|
|
import numpy as np
|
|
import os, shutil
|
|
|
|
import os
|
|
import json
|
|
import math
|
|
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
|
|
from text import text_to_sequence
|
|
|
|
from scipy.io.wavfile import write
|
|
from bertfe import BERTFrontEnd
|
|
from moji.moji import TorchMoji
|
|
|
|
import sys
|
|
|
|
def get_text(text, hps):
|
|
text_norm = text_to_sequence(text, hps.data.text_cleaners)
|
|
if hps.data.add_blank:
|
|
text_norm = commons.intersperse(text_norm, 0)
|
|
text_norm = torch.LongTensor(text_norm)
|
|
return text_norm
|
|
|
|
from text.cleaners import arpa_cleaners
|
|
from text import cleaned_text_to_sequence
|
|
def get_arpa_text(text, hps):
|
|
text_pre = arpa_cleaners(text)
|
|
print(text_pre)
|
|
text_norm = cleaned_text_to_sequence(text_pre)
|
|
if hps.data.add_blank:
|
|
text_norm = commons.intersperse(text_norm, 0)
|
|
text_norm = torch.LongTensor(text_norm)
|
|
return text_norm
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument(
|
|
"--out-path",
|
|
type=str,
|
|
required=True,
|
|
help="Output path of traced model file",
|
|
)
|
|
parser.add_argument(
|
|
"--checkpoint",
|
|
type=str,
|
|
required=True,
|
|
help="Path to model file.",
|
|
)
|
|
parser.add_argument(
|
|
"--config",
|
|
type=str,
|
|
required=True,
|
|
help="Path to model config",
|
|
)
|
|
parser.add_argument(
|
|
"--device",
|
|
type=str,
|
|
required=True,
|
|
help="Device to load model into and export. cpu or cuda",
|
|
)
|
|
parser.add_argument(
|
|
"--test-string",
|
|
type=str,
|
|
required=True,
|
|
help="String to test with",
|
|
)
|
|
parser.add_argument("--bert", default="huawei-noah/TinyBERT_General_4L_312D")
|
|
|
|
args = parser.parse_args()
|
|
|
|
device = torch.device(args.device)
|
|
|
|
if not os.path.isfile(args.checkpoint):
|
|
print(f"Model path {args.checkpoint} incorrect, file not found.")
|
|
|
|
bert_model_name = args.bert
|
|
|
|
print(f"Loading model {args.checkpoint} with config {args.config}...")
|
|
moji = TorchMoji(verbose=True)
|
|
bert_f = BERTFrontEnd(model_name=bert_model_name)
|
|
|
|
hps = utils.get_hparams_from_file(args.config)
|
|
|
|
net_g = SynthesizerTrn(
|
|
len(symbols),
|
|
hps.data.filter_length // 2 + 1,
|
|
hps.train.segment_size // hps.data.hop_length,
|
|
**hps.model).to(device)
|
|
_ = net_g.eval()
|
|
|
|
_ = utils.load_checkpoint(args.checkpoint, net_g, None)
|
|
|
|
test_string = args.test_string
|
|
|
|
if "arpa_cleaners" in hps.data.text_cleaners:
|
|
TT_stn_tst = get_arpa_text(test_string, hps)
|
|
else:
|
|
TT_stn_tst = get_text(test_string, hps)
|
|
|
|
|
|
TT_x_tst = TT_stn_tst.unsqueeze(0).to(device)
|
|
TT_x_tst_lengths = torch.LongTensor([TT_stn_tst.size(0)]).to(device)
|
|
TT_len_scale = torch.FloatTensor([1.]).to(device)
|
|
TT_em = moji(test_string).squeeze().unsqueeze(0).to(device) # [1, 1, 2304] -> [2304] -> [1, 2304]
|
|
|
|
TT_bert_emb, _ = bert_f.infer(test_string) # [1, tokens, channels]
|
|
TT_bert_emb = TT_bert_emb.to(device)
|
|
TT_bert_lens = torch.LongTensor([TT_bert_emb.size(1)]).to(device)
|
|
print("Doing build model run..")
|
|
net_g.infer_ts(TT_x_tst, TT_x_tst_lengths, TT_em, TT_bert_emb, TT_bert_lens,TT_len_scale)
|
|
|
|
|
|
print("Tracing and exporting...")
|
|
traced_vits_noi = torch.jit.trace_module(net_g,
|
|
{"infer_ts_noistft": (TT_x_tst, TT_x_tst_lengths, TT_em, TT_bert_emb, TT_bert_lens,TT_len_scale)}
|
|
,check_trace=False,strict=True)
|
|
|
|
torch.jit.save(traced_vits_noi, args.out_path)
|
|
|
|
|
|
|