File size: 4,963 Bytes
b8db573 |
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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
import os
# os.environ['TORCH_LOGS'] = '+dynamic'
# os.environ['TORCH_LOGS'] = '+export'
# os.environ['TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED']="u0 >= 0"
# os.environ['TORCHDYNAMO_EXTENDED_DEBUG_CPP']="1"
# os.environ['TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL']="u0"
from kokoro import phonemize, tokenize, length_to_mask
import torch.nn.functional as F
from models_scripting import build_model
import torch
from typing import Dict
device = "cpu" #'cuda' if torch.cuda.is_available() else 'cpu'
model = build_model('kokoro-v0_19.pth', device)
voicepack = torch.load('voices/af.pt', weights_only=True).to(device)
speed = 1.
text = "How could I know? It's an unanswerable question. Like asking an unborn child if they'll lead a good life. They haven't even been born."
ps = phonemize(text, "a")
tokens = tokenize(ps)
tokens = torch.LongTensor([[0, *tokens, 0]]).to(device)
class StyleTTS2(torch.nn.Module):
def __init__(self, model, voicepack):
super().__init__()
# self.model = model
self.bert = model.bert
self.bert_encoder = model.bert_encoder
self.predictor = model.predictor
self.decoder = model.decoder
self.text_encoder = model.text_encoder
self.voicepack = voicepack
def forward(self, tokens : torch.Tensor):
speed = 1.
# tokens = torch.nn.functional.pad(tokens, (0, 510 - tokens.shape[-1]))
device = tokens.device
input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)
text_mask = length_to_mask(input_lengths).to(device)
bert_dur = self.bert(tokens)
d_en = self.bert_encoder(bert_dur).transpose(-1, -2)
ref_s = self.voicepack[tokens.shape[1]]
s = ref_s[:, 128:]
d = self.predictor.text_encoder.inference(d_en, s)
x, _ = self.predictor.lstm(d)
duration = self.predictor.duration_proj(x)
duration = torch.sigmoid(duration).sum(axis=-1) / speed
pred_dur = torch.round(duration).clamp(min=1).long()
c_start = F.pad(pred_dur,(1,0), "constant").cumsum(dim=1)[0,0:-1]
c_end = c_start + pred_dur[0,:]
# torch._check(pred_dur.sum().item()>0, lambda: print(f"Got {pred_dur.sum().item()}"))
indices = torch.arange(0, pred_dur.sum().item()).long().to(device)
pred_aln_trg_list=[]
for cs, ce in zip(c_start, c_end):
row = torch.where((indices>=cs) & (indices<ce), 1., 0.)
pred_aln_trg_list.append(row)
pred_aln_trg=torch.vstack(pred_aln_trg_list)
en = d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0).to(device)
F0_pred, N_pred = self.predictor.F0Ntrain(en, s)
t_en = self.text_encoder.inference(tokens)
asr = t_en @ pred_aln_trg.unsqueeze(0).to(device)
return (asr, F0_pred, N_pred, ref_s[:, :128])
# output = self.model.decoder(asr, F0_pred, N_pred, ref_s[:, :128]).squeeze().detach().cpu().numpy()
# bert = torch.jit.script(model.bert)
# bert_encoder = torch.jit.script(model.bert_encoder)
# predictor = torch.jit.script(model.predictor)
# text_encoder = torch.jit.script(model.text_encoder)
# model["bert"] = torch.jit.trace(model["bert"], (tokens, ))
# # model["decoder"] = torch.jit.script(model["decoder"])
# bert_dur = model["bert"](tokens)
# model["bert_encoder"] = torch.jit.trace(model["bert_encoder"], (bert_dur,))
# model["predictor"] = torch.compile(model["predictor"], backend=backend)
# model["text_encoder"] = torch.compile(model["text_encoder"], backend=backend)
style_model = StyleTTS2(model=model, voicepack=voicepack)
style_model.eval()
# style_model = torch.jit.trace_module(style_model.eval(), inputs={'forward': (tokens, )})
# style_model.model["predictor"].F0Ntrain = torch.jit.script(style_model.model["predictor"].F0Ntrain)
(asr, F0_pred, N_pred, ref_s) = style_model(tokens)
print(asr.shape, F0_pred.shape, N_pred.shape, ref_s.shape)
# scripted_style_model = torch.jit.script(style_model)
# (asr, F0_pred, N_pred, ref_s) = scripted_style_model(tokens)
# print(asr.shape, F0_pred.shape, N_pred.shape, ref_s.shape)
# torch.onnx.export(scripted_style_model, ( tokens, ), "style_model.onnx", verbose=True, opset_version=17, input_names=["tokens"], output_names=["asr", "F0_pred", "N_pred", "ref_s"])
# token_len = torch.export.Dim("token_len", min=2, max=510)
# batch = torch.export.Dim("batch")
# dynamic_shapes = {"tokens":{ 1:token_len}}
dynamic_shapes = {"tokens":{ 1:"token_len"}}
print(f"{tokens.shape=}")
torch.onnx.export(model=style_model, args=( tokens, ), dynamic_axes=dynamic_shapes, input_names=["tokens"], f="style_model.onnx",
output_names=["asr", "F0_pred", "N_pred", "ref_s"], opset_version=13, verbose=False, dynamo=False)
# with torch.no_grad():
# torch.export.export(style_model, args=( tokens, ), dynamic_shapes=dynamic_shapes, strict=False)
# export_mod = torch.export.export(style_model, args=( tokens, ), strict=False)
|