File size: 3,283 Bytes
eb932f8 b8db573 eb932f8 |
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 |
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 import build_model
import torch
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)
model = MODEL
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.voicepack = voicepack
def forward(self, tokens):
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.model['bert'](tokens, attention_mask=(~text_mask).int())
d_en = self.model["bert_encoder"](bert_dur).transpose(-1, -2)
ref_s = self.voicepack[tokens.shape[1]]
s = ref_s[:, 128:]
d = self.model["predictor"].text_encoder.inference(d_en, s)
x, _ = self.model["predictor"].lstm(d)
duration = self.model["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.model["predictor"].F0Ntrain(en, s)
t_en = self.model["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()
style_model = StyleTTS2(model=model, voicepack=voicepack)
(asr, F0_pred, N_pred, ref_s) = style_model(tokens)
token_len = torch.export.Dim("token_len", min=2, max=510)
batch = torch.export.Dim("batch")
dynamic_shapes = {"tokens":{0:batch, 1:token_len}}
# with torch.no_grad():
export_mod = torch.export.export(style_model, args=( tokens, ), dynamic_shapes=dynamic_shapes, strict=True)
# export_mod = torch.export.export(style_model, args=( tokens, ), strict=False) |