Added script for testing onnx export.

test.py

@@ -0,0 +1,84 @@
+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=False)
+# export_mod = torch.export.export(style_model, args=( tokens, ), strict=False)