Update ONNXVITS_infer.py

ONNXVITS_infer.py

@@ -125,6 +125,7 @@ class SynthesizerTrn(models.SynthesizerTrn):
                gin_channels=0,
                use_sdp=True,
                emotion_embedding=False,
+               ONNX_dir="./ONNX_net/",
                **kwargs):
 
     super().__init__(
@@ -149,6 +150,7 @@ class SynthesizerTrn(models.SynthesizerTrn):
       use_sdp=use_sdp,
       **kwargs
     )
+    self.ONNX_dir = ONNX_dir
     self.enc_p = TextEncoder(n_vocab,
                              inter_channels,
                              hidden_channels,
@@ -172,7 +174,7 @@ class SynthesizerTrn(models.SynthesizerTrn):
       g = None
 
     # logw = self.dp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w)
-    logw = runonnx("ONNX_net/dp.onnx", x=x.numpy(), x_mask=x_mask.numpy(), g=g.numpy())
+    logw = runonnx(f"{self.ONNX_dir}dp.onnx", x=x.numpy(), x_mask=x_mask.numpy(), g=g.numpy())
     logw = torch.from_numpy(logw[0])
 
     w = torch.exp(logw) * x_mask * length_scale
@@ -189,35 +191,11 @@ class SynthesizerTrn(models.SynthesizerTrn):
     z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
 
     # z = self.flow(z_p, y_mask, g=g, reverse=True)
-    z = runonnx("ONNX_net/flow.onnx", z_p=z_p.numpy(), y_mask=y_mask.numpy(), g=g.numpy())
+    z = runonnx(f"{self.ONNX_dir}flow.onnx", z_p=z_p.numpy(), y_mask=y_mask.numpy(), g=g.numpy())
     z = torch.from_numpy(z[0])
 
     # o = self.dec((z * y_mask)[:,:,:max_len], g=g)
-    o = runonnx("ONNX_net/dec.onnx", z_in=(z * y_mask)[:, :, :max_len].numpy(), g=g.numpy())
+    o = runonnx(f"{self.ONNX_dir}dec.onnx", z_in=(z * y_mask)[:, :, :max_len].numpy(), g=g.numpy())
     o = torch.from_numpy(o[0])
 
-    return o, attn, y_mask, (z, z_p, m_p, logs_p)
-
-  def predict_duration(self, x, x_lengths, sid=None, noise_scale=1, length_scale=1, noise_scale_w=1., max_len=None,
-                       emotion_embedding=None):
-    from ONNXVITS_utils import runonnx
-
-    # x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths)
-    x, m_p, logs_p, x_mask = runonnx("ONNX_net/enc_p.onnx", x=x.numpy(), x_lengths=x_lengths.numpy())
-    x = torch.from_numpy(x)
-    m_p = torch.from_numpy(m_p)
-    logs_p = torch.from_numpy(logs_p)
-    x_mask = torch.from_numpy(x_mask)
-
-    if self.n_speakers > 0:
-      g = self.emb_g(sid).unsqueeze(-1)  # [b, h, 1]
-    else:
-      g = None
-
-    # logw = self.dp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w)
-    logw = runonnx("ONNX_net/dp.onnx", x=x.numpy(), x_mask=x_mask.numpy(), g=g.numpy())
-    logw = torch.from_numpy(logw[0])
-
-    w = torch.exp(logw) * x_mask * length_scale
-    w_ceil = torch.ceil(w)
-    return list(w_ceil.squeeze())
+    return o, attn, y_mask, (z, z_p, m_p, logs_p)