96 lines
3.0 KiB
Python
96 lines
3.0 KiB
Python
from torch import nn
|
|
from infer_pack.models import ( # type:ignore
|
|
TextEncoder256,
|
|
GeneratorNSF,
|
|
PosteriorEncoder,
|
|
ResidualCouplingBlock,
|
|
)
|
|
import torch
|
|
|
|
|
|
class SynthesizerTrnMs256NSFsid_ONNX(nn.Module):
|
|
def __init__(
|
|
self,
|
|
spec_channels,
|
|
segment_size,
|
|
inter_channels,
|
|
hidden_channels,
|
|
filter_channels,
|
|
n_heads,
|
|
n_layers,
|
|
kernel_size,
|
|
p_dropout,
|
|
resblock,
|
|
resblock_kernel_sizes,
|
|
resblock_dilation_sizes,
|
|
upsample_rates,
|
|
upsample_initial_channel,
|
|
upsample_kernel_sizes,
|
|
spk_embed_dim,
|
|
gin_channels,
|
|
sr,
|
|
**kwargs
|
|
):
|
|
super().__init__()
|
|
self.spec_channels = spec_channels
|
|
self.inter_channels = inter_channels
|
|
self.hidden_channels = hidden_channels
|
|
self.filter_channels = filter_channels
|
|
self.n_heads = n_heads
|
|
self.n_layers = n_layers
|
|
self.kernel_size = kernel_size
|
|
self.p_dropout = p_dropout
|
|
self.resblock = resblock
|
|
self.resblock_kernel_sizes = resblock_kernel_sizes
|
|
self.resblock_dilation_sizes = resblock_dilation_sizes
|
|
self.upsample_rates = upsample_rates
|
|
self.upsample_initial_channel = upsample_initial_channel
|
|
self.upsample_kernel_sizes = upsample_kernel_sizes
|
|
self.segment_size = segment_size
|
|
self.gin_channels = gin_channels
|
|
# self.hop_length = hop_length#
|
|
self.spk_embed_dim = spk_embed_dim
|
|
self.enc_p = TextEncoder256(
|
|
inter_channels,
|
|
hidden_channels,
|
|
filter_channels,
|
|
n_heads,
|
|
n_layers,
|
|
kernel_size,
|
|
p_dropout,
|
|
)
|
|
self.dec = GeneratorNSF(
|
|
inter_channels,
|
|
resblock,
|
|
resblock_kernel_sizes,
|
|
resblock_dilation_sizes,
|
|
upsample_rates,
|
|
upsample_initial_channel,
|
|
upsample_kernel_sizes,
|
|
gin_channels=gin_channels,
|
|
sr=sr,
|
|
is_half=kwargs["is_half"],
|
|
)
|
|
self.enc_q = PosteriorEncoder(
|
|
spec_channels,
|
|
inter_channels,
|
|
hidden_channels,
|
|
5,
|
|
1,
|
|
16,
|
|
gin_channels=gin_channels,
|
|
)
|
|
self.flow = ResidualCouplingBlock(
|
|
inter_channels, hidden_channels, 5, 1, 3, gin_channels=gin_channels
|
|
)
|
|
self.emb_g = nn.Embedding(self.spk_embed_dim, gin_channels)
|
|
print("gin_channels:", gin_channels, "self.spk_embed_dim:", self.spk_embed_dim)
|
|
|
|
def forward(self, phone, phone_lengths, pitch, nsff0, sid, max_len=None):
|
|
g = self.emb_g(sid).unsqueeze(-1)
|
|
m_p, logs_p, x_mask = self.enc_p(phone, pitch, phone_lengths)
|
|
z_p = (m_p + torch.exp(logs_p) * torch.randn_like(m_p) * 0.66666) * x_mask
|
|
z = self.flow(z_p, x_mask, g=g, reverse=True)
|
|
o = self.dec((z * x_mask)[:, :, :max_len], nsff0, g=g)
|
|
return o, x_mask, (z, z_p, m_p, logs_p)
|