Refactor: add type hints to models.py / models_jp_extra.py
I didn't add docstring because it is very technical code and I don't understand what is being implemented.
This commit is contained in:
tsukumi
@@ -1,4 +1,5 @@
|
||||
import math
|
||||
from typing import Any, Optional
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
@@ -10,13 +11,18 @@ from style_bert_vits2.models import attentions
|
||||
from style_bert_vits2.models import commons
|
||||
from style_bert_vits2.models import modules
|
||||
from style_bert_vits2.models import monotonic_alignment
|
||||
from style_bert_vits2.nlp.symbols import SYMBOLS, NUM_TONES, NUM_LANGUAGES
|
||||
from style_bert_vits2.nlp.symbols import NUM_LANGUAGES, NUM_TONES, SYMBOLS
|
||||
|
||||
|
||||
class DurationDiscriminator(nn.Module): # vits2
|
||||
def __init__(
|
||||
self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0
|
||||
):
|
||||
self,
|
||||
in_channels: int,
|
||||
filter_channels: int,
|
||||
kernel_size: int,
|
||||
p_dropout: float,
|
||||
gin_channels: int = 0
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
@@ -47,7 +53,7 @@ class DurationDiscriminator(nn.Module): # vits2
|
||||
nn.Linear(2 * filter_channels, 1), nn.Sigmoid()
|
||||
)
|
||||
|
||||
def forward_probability(self, x, dur):
|
||||
def forward_probability(self, x: torch.Tensor, dur: torch.Tensor) -> torch.Tensor:
|
||||
dur = self.dur_proj(dur)
|
||||
x = torch.cat([x, dur], dim=1)
|
||||
x = x.transpose(1, 2)
|
||||
@@ -55,7 +61,14 @@ class DurationDiscriminator(nn.Module): # vits2
|
||||
output_prob = self.output_layer(x)
|
||||
return output_prob
|
||||
|
||||
def forward(self, x, x_mask, dur_r, dur_hat, g=None):
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
x_mask: torch.Tensor,
|
||||
dur_r: torch.Tensor,
|
||||
dur_hat: torch.Tensor,
|
||||
g: Optional[torch.Tensor] = None,
|
||||
) -> list[torch.Tensor]:
|
||||
x = torch.detach(x)
|
||||
if g is not None:
|
||||
g = torch.detach(g)
|
||||
@@ -80,17 +93,17 @@ class DurationDiscriminator(nn.Module): # vits2
|
||||
class TransformerCouplingBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels,
|
||||
hidden_channels,
|
||||
filter_channels,
|
||||
n_heads,
|
||||
n_layers,
|
||||
kernel_size,
|
||||
p_dropout,
|
||||
n_flows=4,
|
||||
gin_channels=0,
|
||||
share_parameter=False,
|
||||
):
|
||||
channels: int,
|
||||
hidden_channels: int,
|
||||
filter_channels: int,
|
||||
n_heads: int,
|
||||
n_layers: int,
|
||||
kernel_size: int,
|
||||
p_dropout: float,
|
||||
n_flows: int = 4,
|
||||
gin_channels: int = 0,
|
||||
share_parameter: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.hidden_channels = hidden_channels
|
||||
@@ -102,16 +115,17 @@ class TransformerCouplingBlock(nn.Module):
|
||||
self.flows = nn.ModuleList()
|
||||
|
||||
self.wn = (
|
||||
attentions.FFT(
|
||||
hidden_channels,
|
||||
filter_channels,
|
||||
n_heads,
|
||||
n_layers,
|
||||
kernel_size,
|
||||
p_dropout,
|
||||
isflow=True,
|
||||
gin_channels=self.gin_channels,
|
||||
)
|
||||
# attentions.FFT(
|
||||
# hidden_channels,
|
||||
# filter_channels,
|
||||
# n_heads,
|
||||
# n_layers,
|
||||
# kernel_size,
|
||||
# p_dropout,
|
||||
# isflow=True,
|
||||
# gin_channels=self.gin_channels,
|
||||
# )
|
||||
None
|
||||
if share_parameter
|
||||
else None
|
||||
)
|
||||
@@ -133,7 +147,13 @@ class TransformerCouplingBlock(nn.Module):
|
||||
)
|
||||
self.flows.append(modules.Flip())
|
||||
|
||||
def forward(self, x, x_mask, g=None, reverse=False):
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
x_mask: torch.Tensor,
|
||||
g: Optional[torch.Tensor] = None,
|
||||
reverse: bool = False,
|
||||
) -> torch.Tensor:
|
||||
if not reverse:
|
||||
for flow in self.flows:
|
||||
x, _ = flow(x, x_mask, g=g, reverse=reverse)
|
||||
@@ -146,13 +166,13 @@ class TransformerCouplingBlock(nn.Module):
|
||||
class StochasticDurationPredictor(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
filter_channels,
|
||||
kernel_size,
|
||||
p_dropout,
|
||||
n_flows=4,
|
||||
gin_channels=0,
|
||||
):
|
||||
in_channels: int,
|
||||
filter_channels: int,
|
||||
kernel_size: int,
|
||||
p_dropout: float,
|
||||
n_flows: int = 4,
|
||||
gin_channels: int = 0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
filter_channels = in_channels # it needs to be removed from future version.
|
||||
self.in_channels = in_channels
|
||||
@@ -192,7 +212,15 @@ class StochasticDurationPredictor(nn.Module):
|
||||
if gin_channels != 0:
|
||||
self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
|
||||
|
||||
def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
x_mask: torch.Tensor,
|
||||
w: Optional[torch.Tensor] = None,
|
||||
g: Optional[torch.Tensor] = None,
|
||||
reverse: bool = False,
|
||||
noise_scale: float = 1.0,
|
||||
) -> torch.Tensor:
|
||||
x = torch.detach(x)
|
||||
x = self.pre(x)
|
||||
if g is not None:
|
||||
@@ -256,8 +284,13 @@ class StochasticDurationPredictor(nn.Module):
|
||||
|
||||
class DurationPredictor(nn.Module):
|
||||
def __init__(
|
||||
self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0
|
||||
):
|
||||
self,
|
||||
in_channels: int,
|
||||
filter_channels: int,
|
||||
kernel_size: int,
|
||||
p_dropout: float,
|
||||
gin_channels: int = 0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
@@ -280,7 +313,7 @@ class DurationPredictor(nn.Module):
|
||||
if gin_channels != 0:
|
||||
self.cond = nn.Conv1d(gin_channels, in_channels, 1)
|
||||
|
||||
def forward(self, x, x_mask, g=None):
|
||||
def forward(self, x: torch.Tensor, x_mask: torch.Tensor, g: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||||
x = torch.detach(x)
|
||||
if g is not None:
|
||||
g = torch.detach(g)
|
||||
@@ -298,14 +331,14 @@ class DurationPredictor(nn.Module):
|
||||
|
||||
|
||||
class Bottleneck(nn.Sequential):
|
||||
def __init__(self, in_dim, hidden_dim):
|
||||
def __init__(self, in_dim: int, hidden_dim: int) -> None:
|
||||
c_fc1 = nn.Linear(in_dim, hidden_dim, bias=False)
|
||||
c_fc2 = nn.Linear(in_dim, hidden_dim, bias=False)
|
||||
super().__init__(*[c_fc1, c_fc2])
|
||||
super().__init__(c_fc1, c_fc2)
|
||||
|
||||
|
||||
class Block(nn.Module):
|
||||
def __init__(self, in_dim, hidden_dim) -> None:
|
||||
def __init__(self, in_dim: int, hidden_dim: int) -> None:
|
||||
super().__init__()
|
||||
self.norm = nn.LayerNorm(in_dim)
|
||||
self.mlp = MLP(in_dim, hidden_dim)
|
||||
@@ -316,13 +349,13 @@ class Block(nn.Module):
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, in_dim, hidden_dim):
|
||||
def __init__(self, in_dim: int, hidden_dim: int) -> None:
|
||||
super().__init__()
|
||||
self.c_fc1 = nn.Linear(in_dim, hidden_dim, bias=False)
|
||||
self.c_fc2 = nn.Linear(in_dim, hidden_dim, bias=False)
|
||||
self.c_proj = nn.Linear(hidden_dim, in_dim, bias=False)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = F.silu(self.c_fc1(x)) * self.c_fc2(x)
|
||||
x = self.c_proj(x)
|
||||
return x
|
||||
@@ -331,16 +364,16 @@ class MLP(nn.Module):
|
||||
class TextEncoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_vocab,
|
||||
out_channels,
|
||||
hidden_channels,
|
||||
filter_channels,
|
||||
n_heads,
|
||||
n_layers,
|
||||
kernel_size,
|
||||
p_dropout,
|
||||
gin_channels=0,
|
||||
):
|
||||
n_vocab: int,
|
||||
out_channels: int,
|
||||
hidden_channels: int,
|
||||
filter_channels: int,
|
||||
n_heads: int,
|
||||
n_layers: int,
|
||||
kernel_size: int,
|
||||
p_dropout: float,
|
||||
gin_channels: int = 0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.n_vocab = n_vocab
|
||||
self.out_channels = out_channels
|
||||
@@ -373,7 +406,16 @@ class TextEncoder(nn.Module):
|
||||
)
|
||||
self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
|
||||
|
||||
def forward(self, x, x_lengths, tone, language, bert, style_vec, g=None):
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
x_lengths: torch.Tensor,
|
||||
tone: torch.Tensor,
|
||||
language: torch.Tensor,
|
||||
bert: torch.Tensor,
|
||||
style_vec: torch.Tensor,
|
||||
g: Optional[torch.Tensor] = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
bert_emb = self.bert_proj(bert).transpose(1, 2)
|
||||
style_emb = self.style_proj(style_vec.unsqueeze(1))
|
||||
x = (
|
||||
@@ -400,14 +442,14 @@ class TextEncoder(nn.Module):
|
||||
class ResidualCouplingBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels,
|
||||
hidden_channels,
|
||||
kernel_size,
|
||||
dilation_rate,
|
||||
n_layers,
|
||||
n_flows=4,
|
||||
gin_channels=0,
|
||||
):
|
||||
channels: int,
|
||||
hidden_channels: int,
|
||||
kernel_size: int,
|
||||
dilation_rate: int,
|
||||
n_layers: int,
|
||||
n_flows: int = 4,
|
||||
gin_channels: int = 0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.hidden_channels = hidden_channels
|
||||
@@ -432,7 +474,13 @@ class ResidualCouplingBlock(nn.Module):
|
||||
)
|
||||
self.flows.append(modules.Flip())
|
||||
|
||||
def forward(self, x, x_mask, g=None, reverse=False):
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
x_mask: torch.Tensor,
|
||||
g: Optional[torch.Tensor] = None,
|
||||
reverse: bool = False,
|
||||
) -> torch.Tensor:
|
||||
if not reverse:
|
||||
for flow in self.flows:
|
||||
x, _ = flow(x, x_mask, g=g, reverse=reverse)
|
||||
@@ -445,14 +493,14 @@ class ResidualCouplingBlock(nn.Module):
|
||||
class PosteriorEncoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
hidden_channels,
|
||||
kernel_size,
|
||||
dilation_rate,
|
||||
n_layers,
|
||||
gin_channels=0,
|
||||
):
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
hidden_channels: int,
|
||||
kernel_size: int,
|
||||
dilation_rate: int,
|
||||
n_layers: int,
|
||||
gin_channels: int = 0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
@@ -472,7 +520,12 @@ class PosteriorEncoder(nn.Module):
|
||||
)
|
||||
self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
|
||||
|
||||
def forward(self, x, x_lengths, g=None):
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
x_lengths: torch.Tensor,
|
||||
g: Optional[torch.Tensor] = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(
|
||||
x.dtype
|
||||
)
|
||||
@@ -487,22 +540,22 @@ class PosteriorEncoder(nn.Module):
|
||||
class Generator(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
initial_channel,
|
||||
resblock,
|
||||
resblock_kernel_sizes,
|
||||
resblock_dilation_sizes,
|
||||
upsample_rates,
|
||||
upsample_initial_channel,
|
||||
upsample_kernel_sizes,
|
||||
gin_channels=0,
|
||||
):
|
||||
initial_channel: int,
|
||||
resblock_str: str,
|
||||
resblock_kernel_sizes: list[int],
|
||||
resblock_dilation_sizes: list[list[int]],
|
||||
upsample_rates: list[int],
|
||||
upsample_initial_channel: int,
|
||||
upsample_kernel_sizes: list[int],
|
||||
gin_channels: int = 0,
|
||||
) -> None:
|
||||
super(Generator, self).__init__()
|
||||
self.num_kernels = len(resblock_kernel_sizes)
|
||||
self.num_upsamples = len(upsample_rates)
|
||||
self.conv_pre = Conv1d(
|
||||
initial_channel, upsample_initial_channel, 7, 1, padding=3
|
||||
)
|
||||
resblock = modules.ResBlock1 if resblock == "1" else modules.ResBlock2
|
||||
resblock = modules.ResBlock1 if resblock_str == "1" else modules.ResBlock2
|
||||
|
||||
self.ups = nn.ModuleList()
|
||||
for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
|
||||
@@ -519,20 +572,22 @@ class Generator(torch.nn.Module):
|
||||
)
|
||||
|
||||
self.resblocks = nn.ModuleList()
|
||||
ch = None
|
||||
for i in range(len(self.ups)):
|
||||
ch = upsample_initial_channel // (2 ** (i + 1))
|
||||
for j, (k, d) in enumerate(
|
||||
zip(resblock_kernel_sizes, resblock_dilation_sizes)
|
||||
):
|
||||
self.resblocks.append(resblock(ch, k, d))
|
||||
self.resblocks.append(resblock(ch, k, d)) # type: ignore
|
||||
|
||||
assert ch is not None
|
||||
self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
|
||||
self.ups.apply(commons.init_weights)
|
||||
|
||||
if gin_channels != 0:
|
||||
self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
|
||||
|
||||
def forward(self, x, g=None):
|
||||
def forward(self, x: torch.Tensor, g: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||||
x = self.conv_pre(x)
|
||||
if g is not None:
|
||||
x = x + self.cond(g)
|
||||
@@ -546,6 +601,7 @@ class Generator(torch.nn.Module):
|
||||
xs = self.resblocks[i * self.num_kernels + j](x)
|
||||
else:
|
||||
xs += self.resblocks[i * self.num_kernels + j](x)
|
||||
assert xs is not None
|
||||
x = xs / self.num_kernels
|
||||
x = F.leaky_relu(x)
|
||||
x = self.conv_post(x)
|
||||
@@ -553,7 +609,7 @@ class Generator(torch.nn.Module):
|
||||
|
||||
return x
|
||||
|
||||
def remove_weight_norm(self):
|
||||
def remove_weight_norm(self) -> None:
|
||||
print("Removing weight norm...")
|
||||
for layer in self.ups:
|
||||
remove_weight_norm(layer)
|
||||
@@ -562,7 +618,7 @@ class Generator(torch.nn.Module):
|
||||
|
||||
|
||||
class DiscriminatorP(torch.nn.Module):
|
||||
def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
|
||||
def __init__(self, period: int, kernel_size: int = 5, stride: int = 3, use_spectral_norm: bool = False) -> None:
|
||||
super(DiscriminatorP, self).__init__()
|
||||
self.period = period
|
||||
self.use_spectral_norm = use_spectral_norm
|
||||
@@ -618,7 +674,7 @@ class DiscriminatorP(torch.nn.Module):
|
||||
)
|
||||
self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
|
||||
|
||||
def forward(self, x):
|
||||
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
|
||||
fmap = []
|
||||
|
||||
# 1d to 2d
|
||||
@@ -641,7 +697,7 @@ class DiscriminatorP(torch.nn.Module):
|
||||
|
||||
|
||||
class DiscriminatorS(torch.nn.Module):
|
||||
def __init__(self, use_spectral_norm=False):
|
||||
def __init__(self, use_spectral_norm: bool = False) -> None:
|
||||
super(DiscriminatorS, self).__init__()
|
||||
norm_f = weight_norm if use_spectral_norm is False else spectral_norm
|
||||
self.convs = nn.ModuleList(
|
||||
@@ -656,7 +712,7 @@ class DiscriminatorS(torch.nn.Module):
|
||||
)
|
||||
self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
|
||||
|
||||
def forward(self, x):
|
||||
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
|
||||
fmap = []
|
||||
|
||||
for layer in self.convs:
|
||||
@@ -671,7 +727,7 @@ class DiscriminatorS(torch.nn.Module):
|
||||
|
||||
|
||||
class MultiPeriodDiscriminator(torch.nn.Module):
|
||||
def __init__(self, use_spectral_norm=False):
|
||||
def __init__(self, use_spectral_norm: bool = False) -> None:
|
||||
super(MultiPeriodDiscriminator, self).__init__()
|
||||
periods = [2, 3, 5, 7, 11]
|
||||
|
||||
@@ -681,7 +737,11 @@ class MultiPeriodDiscriminator(torch.nn.Module):
|
||||
]
|
||||
self.discriminators = nn.ModuleList(discs)
|
||||
|
||||
def forward(self, y, y_hat):
|
||||
def forward(
|
||||
self,
|
||||
y: torch.Tensor,
|
||||
y_hat: torch.Tensor,
|
||||
) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
|
||||
y_d_rs = []
|
||||
y_d_gs = []
|
||||
fmap_rs = []
|
||||
@@ -701,8 +761,12 @@ class WavLMDiscriminator(nn.Module):
|
||||
"""docstring for Discriminator."""
|
||||
|
||||
def __init__(
|
||||
self, slm_hidden=768, slm_layers=13, initial_channel=64, use_spectral_norm=False
|
||||
):
|
||||
self,
|
||||
slm_hidden: int = 768,
|
||||
slm_layers: int = 13,
|
||||
initial_channel: int = 64,
|
||||
use_spectral_norm: bool = False,
|
||||
) -> None:
|
||||
super(WavLMDiscriminator, self).__init__()
|
||||
norm_f = weight_norm if use_spectral_norm == False else spectral_norm
|
||||
self.pre = norm_f(
|
||||
@@ -732,7 +796,7 @@ class WavLMDiscriminator(nn.Module):
|
||||
|
||||
self.conv_post = norm_f(Conv1d(initial_channel * 4, 1, 3, 1, padding=1))
|
||||
|
||||
def forward(self, x):
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.pre(x)
|
||||
|
||||
fmap = []
|
||||
@@ -752,7 +816,7 @@ class ReferenceEncoder(nn.Module):
|
||||
outputs --- [N, ref_enc_gru_size]
|
||||
"""
|
||||
|
||||
def __init__(self, spec_channels, gin_channels=0):
|
||||
def __init__(self, spec_channels: int, gin_channels: int = 0) -> None:
|
||||
super().__init__()
|
||||
self.spec_channels = spec_channels
|
||||
ref_enc_filters = [32, 32, 64, 64, 128, 128]
|
||||
@@ -781,7 +845,7 @@ class ReferenceEncoder(nn.Module):
|
||||
)
|
||||
self.proj = nn.Linear(128, gin_channels)
|
||||
|
||||
def forward(self, inputs, mask=None):
|
||||
def forward(self, inputs: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||||
N = inputs.size(0)
|
||||
out = inputs.view(N, 1, -1, self.spec_channels) # [N, 1, Ty, n_freqs]
|
||||
for conv in self.convs:
|
||||
@@ -799,7 +863,7 @@ class ReferenceEncoder(nn.Module):
|
||||
|
||||
return self.proj(out.squeeze(0))
|
||||
|
||||
def calculate_channels(self, L, kernel_size, stride, pad, n_convs):
|
||||
def calculate_channels(self, L: int, kernel_size: int, stride: int, pad: int, n_convs: int) -> int:
|
||||
for i in range(n_convs):
|
||||
L = (L - kernel_size + 2 * pad) // stride + 1
|
||||
return L
|
||||
@@ -812,31 +876,31 @@ class SynthesizerTrn(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
n_vocab,
|
||||
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,
|
||||
n_speakers=256,
|
||||
gin_channels=256,
|
||||
use_sdp=True,
|
||||
n_flow_layer=4,
|
||||
n_layers_trans_flow=6,
|
||||
flow_share_parameter=False,
|
||||
use_transformer_flow=True,
|
||||
**kwargs
|
||||
):
|
||||
n_vocab: int,
|
||||
spec_channels: int,
|
||||
segment_size: int,
|
||||
inter_channels: int,
|
||||
hidden_channels: int,
|
||||
filter_channels: int,
|
||||
n_heads: int,
|
||||
n_layers: int,
|
||||
kernel_size: int,
|
||||
p_dropout: float,
|
||||
resblock: str,
|
||||
resblock_kernel_sizes: list[int],
|
||||
resblock_dilation_sizes: list[list[int]],
|
||||
upsample_rates: list[int],
|
||||
upsample_initial_channel: int,
|
||||
upsample_kernel_sizes: list[int],
|
||||
n_speakers: int = 256,
|
||||
gin_channels: int = 256,
|
||||
use_sdp: bool = True,
|
||||
n_flow_layer: int = 4,
|
||||
n_layers_trans_flow: int = 6,
|
||||
flow_share_parameter: bool = False,
|
||||
use_transformer_flow: bool = True,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.n_vocab = n_vocab
|
||||
self.spec_channels = spec_channels
|
||||
@@ -933,16 +997,26 @@ class SynthesizerTrn(nn.Module):
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
x_lengths,
|
||||
y,
|
||||
y_lengths,
|
||||
sid,
|
||||
tone,
|
||||
language,
|
||||
bert,
|
||||
style_vec,
|
||||
):
|
||||
x: torch.Tensor,
|
||||
x_lengths: torch.Tensor,
|
||||
y: torch.Tensor,
|
||||
y_lengths: torch.Tensor,
|
||||
sid: torch.Tensor,
|
||||
tone: torch.Tensor,
|
||||
language: torch.Tensor,
|
||||
bert: torch.Tensor,
|
||||
style_vec: torch.Tensor,
|
||||
) -> tuple[
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
tuple[torch.Tensor, ...],
|
||||
tuple[torch.Tensor, ...],
|
||||
]:
|
||||
if self.n_speakers > 0:
|
||||
g = self.emb_g(sid).unsqueeze(-1) # [b, h, 1]
|
||||
else:
|
||||
@@ -1014,32 +1088,33 @@ class SynthesizerTrn(nn.Module):
|
||||
ids_slice,
|
||||
x_mask,
|
||||
y_mask,
|
||||
(z, z_p, m_p, logs_p, m_q, logs_q),
|
||||
(z, z_p, m_p, logs_p, m_q, logs_q), # type: ignore
|
||||
(x, logw, logw_), # , logw_sdp),
|
||||
g,
|
||||
)
|
||||
|
||||
def infer(
|
||||
self,
|
||||
x,
|
||||
x_lengths,
|
||||
sid,
|
||||
tone,
|
||||
language,
|
||||
bert,
|
||||
style_vec,
|
||||
noise_scale=0.667,
|
||||
length_scale=1.0,
|
||||
noise_scale_w=0.8,
|
||||
max_len=None,
|
||||
sdp_ratio=0.0,
|
||||
y=None,
|
||||
):
|
||||
x: torch.Tensor,
|
||||
x_lengths: torch.Tensor,
|
||||
sid: torch.Tensor,
|
||||
tone: torch.Tensor,
|
||||
language: torch.Tensor,
|
||||
bert: torch.Tensor,
|
||||
style_vec: torch.Tensor,
|
||||
noise_scale: float = 0.667,
|
||||
length_scale: float = 1.0,
|
||||
noise_scale_w: float = 0.8,
|
||||
max_len: Optional[int] = None,
|
||||
sdp_ratio: float = 0.0,
|
||||
y: Optional[torch.Tensor] = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor, ...]]:
|
||||
# x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths, tone, language, bert)
|
||||
# g = self.gst(y)
|
||||
if self.n_speakers > 0:
|
||||
g = self.emb_g(sid).unsqueeze(-1) # [b, h, 1]
|
||||
else:
|
||||
assert y is not None
|
||||
g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
|
||||
x, m_p, logs_p, x_mask = self.enc_p(
|
||||
x, x_lengths, tone, language, bert, style_vec, g=g
|
||||
|
||||
Reference in New Issue
Block a user