tuna2134/sbv2-v2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Refactor: add type hints to models.py / models_jp_extra.py

Browse Source 
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
2024-03-08 23:52:44 +00:00
parent e1fad54c99
commit 8feef04cef
4 changed files with 490 additions and 325 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
style_bert_vits2/models/attentions.py
View File

@@ -9,7 +9,7 @@ from style_bert_vits2.models import commons
class LayerNorm(nn.Module): class LayerNorm(nn.Module):
def __init__(self, channels: int, eps: float = 1e-5): def __init__(self, channels: int, eps: float = 1e-5) -> None:
super().__init__() super().__init__()
self.channels = channels self.channels = channels
self.eps = eps self.eps = eps
@@ -45,7 +45,7 @@ class Encoder(nn.Module):
window_size: int = 4, window_size: int = 4,
isflow: bool = True, isflow: bool = True,
**kwargs: Any **kwargs: Any
): ) -> None:
super().__init__() super().__init__()
self.hidden_channels = hidden_channels self.hidden_channels = hidden_channels
self.filter_channels = filter_channels self.filter_channels = filter_channels
@@ -132,7 +132,7 @@ class Decoder(nn.Module):
proximal_bias: bool = False, proximal_bias: bool = False,
proximal_init: bool = True, proximal_init: bool = True,
**kwargs: Any **kwargs: Any
): ) -> None:
super().__init__() super().__init__()
self.hidden_channels = hidden_channels self.hidden_channels = hidden_channels
self.filter_channels = filter_channels self.filter_channels = filter_channels
@@ -180,7 +180,7 @@ class Decoder(nn.Module):
) )
self.norm_layers_2.append(LayerNorm(hidden_channels)) self.norm_layers_2.append(LayerNorm(hidden_channels))
def forward(self, x: torch.Tensor, x_mask: torch.Tensor, h: torch.Tensor, h_mask: torch.Tensor): def forward(self, x: torch.Tensor, x_mask: torch.Tensor, h: torch.Tensor, h_mask: torch.Tensor) -> torch.Tensor:
""" """
x: decoder input x: decoder input
h: encoder output h: encoder output
@@ -218,7 +218,7 @@ class MultiHeadAttention(nn.Module):
block_length: Optional[int] = None, block_length: Optional[int] = None,
proximal_bias: bool = False, proximal_bias: bool = False,
proximal_init: bool = False, proximal_init: bool = False,
): ) -> None:
super().__init__() super().__init__()
assert channels % n_heads == 0 assert channels % n_heads == 0
@@ -272,7 +272,13 @@ class MultiHeadAttention(nn.Module):
x = self.conv_o(x) x = self.conv_o(x)
return x return x
def attention(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: Optional[torch.Tensor] = None) -> tuple[torch.Tensor, torch.Tensor]: def attention(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
mask: Optional[torch.Tensor] = None,
) -> tuple[torch.Tensor, torch.Tensor]:
# reshape [b, d, t] -> [b, n_h, t, d_k] # reshape [b, d, t] -> [b, n_h, t, d_k]
b, d, t_s, t_t = (*key.size(), query.size(2)) b, d, t_s, t_t = (*key.size(), query.size(2))
query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3) query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
@@ -419,7 +425,7 @@ class FFN(nn.Module):
p_dropout: float = 0.0, p_dropout: float = 0.0,
activation: Optional[str] = None, activation: Optional[str] = None,
causal: bool = False, causal: bool = False,
): ) -> None:
super().__init__() super().__init__()
self.in_channels = in_channels self.in_channels = in_channels
self.out_channels = out_channels self.out_channels = out_channels

384
style_bert_vits2/models/models.py
View File

@@ -1,4 +1,5 @@
import math import math
from typing import Any, Optional
import torch import torch
from torch import nn from torch import nn
@@ -10,14 +11,18 @@ from style_bert_vits2.models import attentions
from style_bert_vits2.models import commons from style_bert_vits2.models import commons
from style_bert_vits2.models import modules from style_bert_vits2.models import modules
from style_bert_vits2.models import monotonic_alignment from style_bert_vits2.models import monotonic_alignment
from style_bert_vits2.models.commons import get_padding, init_weights
from style_bert_vits2.nlp.symbols import NUM_LANGUAGES, NUM_TONES, SYMBOLS from style_bert_vits2.nlp.symbols import NUM_LANGUAGES, NUM_TONES, SYMBOLS
class DurationDiscriminator(nn.Module): # vits2 class DurationDiscriminator(nn.Module): # vits2
def __init__( 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__() super().__init__()
self.in_channels = in_channels self.in_channels = in_channels
@@ -51,7 +56,13 @@ class DurationDiscriminator(nn.Module): # vits2
self.output_layer = nn.Sequential(nn.Linear(filter_channels, 1), nn.Sigmoid()) self.output_layer = nn.Sequential(nn.Linear(filter_channels, 1), nn.Sigmoid())
def forward_probability(self, x, x_mask, dur, g=None): def forward_probability(
self,
x: torch.Tensor,
x_mask: torch.Tensor,
dur: torch.Tensor,
g: Optional[torch.Tensor] = None,
) -> torch.Tensor:
dur = self.dur_proj(dur) dur = self.dur_proj(dur)
x = torch.cat([x, dur], dim=1) x = torch.cat([x, dur], dim=1)
x = self.pre_out_conv_1(x * x_mask) x = self.pre_out_conv_1(x * x_mask)
@@ -67,7 +78,14 @@ class DurationDiscriminator(nn.Module): # vits2
output_prob = self.output_layer(x) output_prob = self.output_layer(x)
return output_prob 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) x = torch.detach(x)
if g is not None: if g is not None:
g = torch.detach(g) g = torch.detach(g)
@@ -92,17 +110,17 @@ class DurationDiscriminator(nn.Module): # vits2
class TransformerCouplingBlock(nn.Module): class TransformerCouplingBlock(nn.Module):
def __init__( def __init__(
self, self,
channels, channels: int,
hidden_channels, hidden_channels: int,
filter_channels, filter_channels: int,
n_heads, n_heads: int,
n_layers, n_layers: int,
kernel_size, kernel_size: int,
p_dropout, p_dropout: float,
n_flows=4, n_flows: int = 4,
gin_channels=0, gin_channels: int = 0,
share_parameter=False, share_parameter: bool = False,
): ) -> None:
super().__init__() super().__init__()
self.channels = channels self.channels = channels
self.hidden_channels = hidden_channels self.hidden_channels = hidden_channels
@@ -114,16 +132,17 @@ class TransformerCouplingBlock(nn.Module):
self.flows = nn.ModuleList() self.flows = nn.ModuleList()
self.wn = ( self.wn = (
attentions.FFT( # attentions.FFT(
hidden_channels, # hidden_channels,
filter_channels, # filter_channels,
n_heads, # n_heads,
n_layers, # n_layers,
kernel_size, # kernel_size,
p_dropout, # p_dropout,
isflow=True, # isflow=True,
gin_channels=self.gin_channels, # gin_channels=self.gin_channels,
) # )
None
if share_parameter if share_parameter
else None else None
) )
@@ -145,7 +164,13 @@ class TransformerCouplingBlock(nn.Module):
) )
self.flows.append(modules.Flip()) 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: if not reverse:
for flow in self.flows: for flow in self.flows:
x, _ = flow(x, x_mask, g=g, reverse=reverse) x, _ = flow(x, x_mask, g=g, reverse=reverse)
@@ -158,13 +183,13 @@ class TransformerCouplingBlock(nn.Module):
class StochasticDurationPredictor(nn.Module): class StochasticDurationPredictor(nn.Module):
def __init__( def __init__(
self, self,
in_channels, in_channels: int,
filter_channels, filter_channels: int,
kernel_size, kernel_size: int,
p_dropout, p_dropout: float,
n_flows=4, n_flows: int = 4,
gin_channels=0, gin_channels: int = 0,
): ) -> None:
super().__init__() super().__init__()
filter_channels = in_channels # it needs to be removed from future version. filter_channels = in_channels # it needs to be removed from future version.
self.in_channels = in_channels self.in_channels = in_channels
@@ -204,7 +229,15 @@ class StochasticDurationPredictor(nn.Module):
if gin_channels != 0: if gin_channels != 0:
self.cond = nn.Conv1d(gin_channels, filter_channels, 1) 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 = torch.detach(x)
x = self.pre(x) x = self.pre(x)
if g is not None: if g is not None:
@@ -268,8 +301,13 @@ class StochasticDurationPredictor(nn.Module):
class DurationPredictor(nn.Module): class DurationPredictor(nn.Module):
def __init__( 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__() super().__init__()
self.in_channels = in_channels self.in_channels = in_channels
@@ -292,7 +330,7 @@ class DurationPredictor(nn.Module):
if gin_channels != 0: if gin_channels != 0:
self.cond = nn.Conv1d(gin_channels, in_channels, 1) 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) x = torch.detach(x)
if g is not None: if g is not None:
g = torch.detach(g) g = torch.detach(g)
@@ -312,17 +350,17 @@ class DurationPredictor(nn.Module):
class TextEncoder(nn.Module): class TextEncoder(nn.Module):
def __init__( def __init__(
self, self,
n_vocab, n_vocab: int,
out_channels, out_channels: int,
hidden_channels, hidden_channels: int,
filter_channels, filter_channels: int,
n_heads, n_heads: int,
n_layers, n_layers: int,
kernel_size, kernel_size: int,
p_dropout, p_dropout: float,
n_speakers, n_speakers: int,
gin_channels=0, gin_channels: int = 0,
): ) -> None:
super().__init__() super().__init__()
self.n_vocab = n_vocab self.n_vocab = n_vocab
self.out_channels = out_channels self.out_channels = out_channels
@@ -357,17 +395,17 @@ class TextEncoder(nn.Module):
def forward( def forward(
self, self,
x, x: torch.Tensor,
x_lengths, x_lengths: torch.Tensor,
tone, tone: torch.Tensor,
language, language: torch.Tensor,
bert, bert: torch.Tensor,
ja_bert, ja_bert: torch.Tensor,
en_bert, en_bert: torch.Tensor,
style_vec, style_vec: torch.Tensor,
sid, sid: torch.Tensor,
g=None, g: Optional[torch.Tensor] = None,
): ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
bert_emb = self.bert_proj(bert).transpose(1, 2) bert_emb = self.bert_proj(bert).transpose(1, 2)
ja_bert_emb = self.ja_bert_proj(ja_bert).transpose(1, 2) ja_bert_emb = self.ja_bert_proj(ja_bert).transpose(1, 2)
en_bert_emb = self.en_bert_proj(en_bert).transpose(1, 2) en_bert_emb = self.en_bert_proj(en_bert).transpose(1, 2)
@@ -399,14 +437,14 @@ class TextEncoder(nn.Module):
class ResidualCouplingBlock(nn.Module): class ResidualCouplingBlock(nn.Module):
def __init__( def __init__(
self, self,
channels, channels: int,
hidden_channels, hidden_channels: int,
kernel_size, kernel_size: int,
dilation_rate, dilation_rate: int,
n_layers, n_layers: int,
n_flows=4, n_flows: int = 4,
gin_channels=0, gin_channels: int = 0,
): ) -> None:
super().__init__() super().__init__()
self.channels = channels self.channels = channels
self.hidden_channels = hidden_channels self.hidden_channels = hidden_channels
@@ -431,7 +469,13 @@ class ResidualCouplingBlock(nn.Module):
) )
self.flows.append(modules.Flip()) 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: if not reverse:
for flow in self.flows: for flow in self.flows:
x, _ = flow(x, x_mask, g=g, reverse=reverse) x, _ = flow(x, x_mask, g=g, reverse=reverse)
@@ -444,14 +488,14 @@ class ResidualCouplingBlock(nn.Module):
class PosteriorEncoder(nn.Module): class PosteriorEncoder(nn.Module):
def __init__( def __init__(
self, self,
in_channels, in_channels: int,
out_channels, out_channels: int,
hidden_channels, hidden_channels: int,
kernel_size, kernel_size: int,
dilation_rate, dilation_rate: int,
n_layers, n_layers: int,
gin_channels=0, gin_channels: int = 0,
): ) -> None:
super().__init__() super().__init__()
self.in_channels = in_channels self.in_channels = in_channels
self.out_channels = out_channels self.out_channels = out_channels
@@ -471,7 +515,12 @@ class PosteriorEncoder(nn.Module):
) )
self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1) 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_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(
x.dtype x.dtype
) )
@@ -486,22 +535,22 @@ class PosteriorEncoder(nn.Module):
class Generator(torch.nn.Module): class Generator(torch.nn.Module):
def __init__( def __init__(
self, self,
initial_channel, initial_channel: int,
resblock, resblock_str: str,
resblock_kernel_sizes, resblock_kernel_sizes: list[int],
resblock_dilation_sizes, resblock_dilation_sizes: list[list[int]],
upsample_rates, upsample_rates: list[int],
upsample_initial_channel, upsample_initial_channel: int,
upsample_kernel_sizes, upsample_kernel_sizes: list[int],
gin_channels=0, gin_channels: int = 0,
): ) -> None:
super(Generator, self).__init__() super(Generator, self).__init__()
self.num_kernels = len(resblock_kernel_sizes) self.num_kernels = len(resblock_kernel_sizes)
self.num_upsamples = len(upsample_rates) self.num_upsamples = len(upsample_rates)
self.conv_pre = Conv1d( self.conv_pre = Conv1d(
initial_channel, upsample_initial_channel, 7, 1, padding=3 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() self.ups = nn.ModuleList()
for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
@@ -518,20 +567,22 @@ class Generator(torch.nn.Module):
) )
self.resblocks = nn.ModuleList() self.resblocks = nn.ModuleList()
ch = None
for i in range(len(self.ups)): for i in range(len(self.ups)):
ch = upsample_initial_channel // (2 ** (i + 1)) ch = upsample_initial_channel // (2 ** (i + 1))
for j, (k, d) in enumerate( for j, (k, d) in enumerate(
zip(resblock_kernel_sizes, resblock_dilation_sizes) 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.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
self.ups.apply(init_weights) self.ups.apply(commons.init_weights)
if gin_channels != 0: if gin_channels != 0:
self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1) 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) x = self.conv_pre(x)
if g is not None: if g is not None:
x = x + self.cond(g) x = x + self.cond(g)
@@ -545,6 +596,7 @@ class Generator(torch.nn.Module):
xs = self.resblocks[i * self.num_kernels + j](x) xs = self.resblocks[i * self.num_kernels + j](x)
else: else:
xs += self.resblocks[i * self.num_kernels + j](x) xs += self.resblocks[i * self.num_kernels + j](x)
assert xs is not None
x = xs / self.num_kernels x = xs / self.num_kernels
x = F.leaky_relu(x) x = F.leaky_relu(x)
x = self.conv_post(x) x = self.conv_post(x)
@@ -552,7 +604,7 @@ class Generator(torch.nn.Module):
return x return x
def remove_weight_norm(self): def remove_weight_norm(self) -> None:
print("Removing weight norm...") print("Removing weight norm...")
for layer in self.ups: for layer in self.ups:
remove_weight_norm(layer) remove_weight_norm(layer)
@@ -561,7 +613,7 @@ class Generator(torch.nn.Module):
class DiscriminatorP(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__() super(DiscriminatorP, self).__init__()
self.period = period self.period = period
self.use_spectral_norm = use_spectral_norm self.use_spectral_norm = use_spectral_norm
@@ -574,7 +626,7 @@ class DiscriminatorP(torch.nn.Module):
32, 32,
(kernel_size, 1), (kernel_size, 1),
(stride, 1), (stride, 1),
padding=(get_padding(kernel_size, 1), 0), padding=(commons.get_padding(kernel_size, 1), 0),
) )
), ),
norm_f( norm_f(
@@ -583,7 +635,7 @@ class DiscriminatorP(torch.nn.Module):
128, 128,
(kernel_size, 1), (kernel_size, 1),
(stride, 1), (stride, 1),
padding=(get_padding(kernel_size, 1), 0), padding=(commons.get_padding(kernel_size, 1), 0),
) )
), ),
norm_f( norm_f(
@@ -592,7 +644,7 @@ class DiscriminatorP(torch.nn.Module):
512, 512,
(kernel_size, 1), (kernel_size, 1),
(stride, 1), (stride, 1),
padding=(get_padding(kernel_size, 1), 0), padding=(commons.get_padding(kernel_size, 1), 0),
) )
), ),
norm_f( norm_f(
@@ -601,7 +653,7 @@ class DiscriminatorP(torch.nn.Module):
1024, 1024,
(kernel_size, 1), (kernel_size, 1),
(stride, 1), (stride, 1),
padding=(get_padding(kernel_size, 1), 0), padding=(commons.get_padding(kernel_size, 1), 0),
) )
), ),
norm_f( norm_f(
@@ -610,14 +662,14 @@ class DiscriminatorP(torch.nn.Module):
1024, 1024,
(kernel_size, 1), (kernel_size, 1),
1, 1,
padding=(get_padding(kernel_size, 1), 0), padding=(commons.get_padding(kernel_size, 1), 0),
) )
), ),
] ]
) )
self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0))) 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 = [] fmap = []
# 1d to 2d # 1d to 2d
@@ -640,7 +692,7 @@ class DiscriminatorP(torch.nn.Module):
class DiscriminatorS(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__() super(DiscriminatorS, self).__init__()
norm_f = weight_norm if use_spectral_norm is False else spectral_norm norm_f = weight_norm if use_spectral_norm is False else spectral_norm
self.convs = nn.ModuleList( self.convs = nn.ModuleList(
@@ -655,7 +707,7 @@ class DiscriminatorS(torch.nn.Module):
) )
self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1)) 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 = [] fmap = []
for layer in self.convs: for layer in self.convs:
@@ -670,7 +722,7 @@ class DiscriminatorS(torch.nn.Module):
class MultiPeriodDiscriminator(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__() super(MultiPeriodDiscriminator, self).__init__()
periods = [2, 3, 5, 7, 11] periods = [2, 3, 5, 7, 11]
@@ -680,7 +732,11 @@ class MultiPeriodDiscriminator(torch.nn.Module):
] ]
self.discriminators = nn.ModuleList(discs) 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_rs = []
y_d_gs = [] y_d_gs = []
fmap_rs = [] fmap_rs = []
@@ -702,7 +758,7 @@ class ReferenceEncoder(nn.Module):
outputs --- [N, ref_enc_gru_size] 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__() super().__init__()
self.spec_channels = spec_channels self.spec_channels = spec_channels
ref_enc_filters = [32, 32, 64, 64, 128, 128] ref_enc_filters = [32, 32, 64, 64, 128, 128]
@@ -731,7 +787,7 @@ class ReferenceEncoder(nn.Module):
) )
self.proj = nn.Linear(128, gin_channels) 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) N = inputs.size(0)
out = inputs.view(N, 1, -1, self.spec_channels) # [N, 1, Ty, n_freqs] out = inputs.view(N, 1, -1, self.spec_channels) # [N, 1, Ty, n_freqs]
for conv in self.convs: for conv in self.convs:
@@ -749,7 +805,7 @@ class ReferenceEncoder(nn.Module):
return self.proj(out.squeeze(0)) 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): for i in range(n_convs):
L = (L - kernel_size + 2 * pad) // stride + 1 L = (L - kernel_size + 2 * pad) // stride + 1
return L return L
@@ -762,31 +818,31 @@ class SynthesizerTrn(nn.Module):
def __init__( def __init__(
self, self,
n_vocab, n_vocab: int,
spec_channels, spec_channels: int,
segment_size, segment_size: int,
inter_channels, inter_channels: int,
hidden_channels, hidden_channels: int,
filter_channels, filter_channels: int,
n_heads, n_heads: int,
n_layers, n_layers: int,
kernel_size, kernel_size: int,
p_dropout, p_dropout: float,
resblock, resblock: str,
resblock_kernel_sizes, resblock_kernel_sizes: list[int],
resblock_dilation_sizes, resblock_dilation_sizes: list[list[int]],
upsample_rates, upsample_rates: list[int],
upsample_initial_channel, upsample_initial_channel: int,
upsample_kernel_sizes, upsample_kernel_sizes: list[int],
n_speakers=256, n_speakers: int = 256,
gin_channels=256, gin_channels: int = 256,
use_sdp=True, use_sdp: bool = True,
n_flow_layer=4, n_flow_layer: int = 4,
n_layers_trans_flow=4, n_layers_trans_flow: int = 4,
flow_share_parameter=False, flow_share_parameter: bool = False,
use_transformer_flow=True, use_transformer_flow: bool = True,
**kwargs, **kwargs: Any,
): ) -> None:
super().__init__() super().__init__()
self.n_vocab = n_vocab self.n_vocab = n_vocab
self.spec_channels = spec_channels self.spec_channels = spec_channels
@@ -884,18 +940,27 @@ class SynthesizerTrn(nn.Module):
def forward( def forward(
self, self,
x, x: torch.Tensor,
x_lengths, x_lengths: torch.Tensor,
y, y: torch.Tensor,
y_lengths, y_lengths: torch.Tensor,
sid, sid: torch.Tensor,
tone, tone: torch.Tensor,
language, language: torch.Tensor,
bert, bert: torch.Tensor,
ja_bert, ja_bert: torch.Tensor,
en_bert, en_bert: torch.Tensor,
style_vec, style_vec: torch.Tensor,
): ) -> tuple[
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
tuple[torch.Tensor, ...],
tuple[torch.Tensor, ...],
]:
if self.n_speakers > 0: if self.n_speakers > 0:
g = self.emb_g(sid).unsqueeze(-1) # [b, h, 1] g = self.emb_g(sid).unsqueeze(-1) # [b, h, 1]
else: else:
@@ -973,27 +1038,28 @@ class SynthesizerTrn(nn.Module):
def infer( def infer(
self, self,
x, x: torch.Tensor,
x_lengths, x_lengths: torch.Tensor,
sid, sid: torch.Tensor,
tone, tone: torch.Tensor,
language, language: torch.Tensor,
bert, bert: torch.Tensor,
ja_bert, ja_bert: torch.Tensor,
en_bert, en_bert: torch.Tensor,
style_vec, style_vec: torch.Tensor,
noise_scale=0.667, noise_scale: float = 0.667,
length_scale=1.0, length_scale: float = 1.0,
noise_scale_w=0.8, noise_scale_w: float = 0.8,
max_len=None, max_len: Optional[int] = None,
sdp_ratio=0.0, sdp_ratio: float = 0.0,
y=None, 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) # x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths, tone, language, bert)
# g = self.gst(y) # g = self.gst(y)
if self.n_speakers > 0: if self.n_speakers > 0:
g = self.emb_g(sid).unsqueeze(-1) # [b, h, 1] g = self.emb_g(sid).unsqueeze(-1) # [b, h, 1]
else: else:
assert y is not None
g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1) g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
x, m_p, logs_p, x_mask = self.enc_p( x, m_p, logs_p, x_mask = self.enc_p(
x, x_lengths, tone, language, bert, ja_bert, en_bert, style_vec, sid, g=g x, x_lengths, tone, language, bert, ja_bert, en_bert, style_vec, sid, g=g

369
style_bert_vits2/models/models_jp_extra.py
View File

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

42
style_bert_vits2/models/modules.py
View File

@@ -16,7 +16,7 @@ LRELU_SLOPE = 0.1
class LayerNorm(nn.Module): class LayerNorm(nn.Module):
def __init__(self, channels: int, eps: float = 1e-5): def __init__(self, channels: int, eps: float = 1e-5) -> None:
super().__init__() super().__init__()
self.channels = channels self.channels = channels
self.eps = eps self.eps = eps
@@ -39,7 +39,7 @@ class ConvReluNorm(nn.Module):
kernel_size: int, kernel_size: int,
n_layers: int, n_layers: int,
p_dropout: float, p_dropout: float,
): ) -> None:
super().__init__() super().__init__()
self.in_channels = in_channels self.in_channels = in_channels
self.hidden_channels = hidden_channels self.hidden_channels = hidden_channels
@@ -88,7 +88,7 @@ class DDSConv(nn.Module):
Dialted and Depth-Separable Convolution Dialted and Depth-Separable Convolution
""" """
def __init__(self, channels: int, kernel_size: int, n_layers: int, p_dropout: float = 0.0): def __init__(self, channels: int, kernel_size: int, n_layers: int, p_dropout: float = 0.0) -> None:
super().__init__() super().__init__()
self.channels = channels self.channels = channels
self.kernel_size = kernel_size self.kernel_size = kernel_size
@@ -141,7 +141,7 @@ class WN(torch.nn.Module):
n_layers: int, n_layers: int,
gin_channels: int = 0, gin_channels: int = 0,
p_dropout: float = 0, p_dropout: float = 0,
): ) -> None:
super(WN, self).__init__() super(WN, self).__init__()
assert kernel_size % 2 == 1 assert kernel_size % 2 == 1
self.hidden_channels = hidden_channels self.hidden_channels = hidden_channels
@@ -221,7 +221,7 @@ class WN(torch.nn.Module):
class ResBlock1(torch.nn.Module): class ResBlock1(torch.nn.Module):
def __init__(self, channels: int, kernel_size: int = 3, dilation: tuple[int, int, int] = (1, 3, 5)): def __init__(self, channels: int, kernel_size: int = 3, dilation: tuple[int, int, int] = (1, 3, 5)) -> None:
super(ResBlock1, self).__init__() super(ResBlock1, self).__init__()
self.convs1 = nn.ModuleList( self.convs1 = nn.ModuleList(
[ [
@@ -318,7 +318,7 @@ class ResBlock1(torch.nn.Module):
class ResBlock2(torch.nn.Module): class ResBlock2(torch.nn.Module):
def __init__(self, channels: int, kernel_size: int = 3, dilation: tuple[int, int] = (1, 3)): def __init__(self, channels: int, kernel_size: int = 3, dilation: tuple[int, int] = (1, 3)) -> None:
super(ResBlock2, self).__init__() super(ResBlock2, self).__init__()
self.convs = nn.ModuleList( self.convs = nn.ModuleList(
[ [
@@ -363,7 +363,13 @@ class ResBlock2(torch.nn.Module):
class Log(nn.Module): class Log(nn.Module):
def forward(self, x: torch.Tensor, x_mask: torch.Tensor, reverse: bool = False, **kwargs: Any): def forward(
self,
x: torch.Tensor,
x_mask: torch.Tensor,
reverse: bool = False,
**kwargs: Any,
) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
if not reverse: if not reverse:
y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
logdet = torch.sum(-y, [1, 2]) logdet = torch.sum(-y, [1, 2])
@@ -390,7 +396,7 @@ class Flip(nn.Module):
class ElementwiseAffine(nn.Module): class ElementwiseAffine(nn.Module):
def __init__(self, channels: int): def __init__(self, channels: int) -> None:
super().__init__() super().__init__()
self.channels = channels self.channels = channels
self.m = nn.Parameter(torch.zeros(channels, 1)) self.m = nn.Parameter(torch.zeros(channels, 1))
@@ -424,7 +430,7 @@ class ResidualCouplingLayer(nn.Module):
p_dropout: float = 0, p_dropout: float = 0,
gin_channels: int = 0, gin_channels: int = 0,
mean_only: bool = False, mean_only: bool = False,
): ) -> None:
assert channels % 2 == 0, "channels should be divisible by 2" assert channels % 2 == 0, "channels should be divisible by 2"
super().__init__() super().__init__()
self.channels = channels self.channels = channels
@@ -449,7 +455,13 @@ class ResidualCouplingLayer(nn.Module):
assert self.post.bias is not None assert self.post.bias is not None
self.post.bias.data.zero_() self.post.bias.data.zero_()
def forward(self, x: torch.Tensor, x_mask: torch.Tensor, g: Optional[torch.Tensor] = None, reverse: bool = False): def forward(
self,
x: torch.Tensor,
x_mask: torch.Tensor,
g: Optional[torch.Tensor] = None,
reverse: bool = False,
) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
x0, x1 = torch.split(x, [self.half_channels] * 2, 1) x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
h = self.pre(x0) * x_mask h = self.pre(x0) * x_mask
h = self.enc(h, x_mask, g=g) h = self.enc(h, x_mask, g=g)
@@ -480,7 +492,7 @@ class ConvFlow(nn.Module):
n_layers: int, n_layers: int,
num_bins: int = 10, num_bins: int = 10,
tail_bound: float = 5.0, tail_bound: float = 5.0,
): ) -> None:
super().__init__() super().__init__()
self.in_channels = in_channels self.in_channels = in_channels
self.filter_channels = filter_channels self.filter_channels = filter_channels
@@ -499,7 +511,13 @@ class ConvFlow(nn.Module):
assert self.proj.bias is not None assert self.proj.bias is not None
self.proj.bias.data.zero_() self.proj.bias.data.zero_()
def forward(self, x: torch.Tensor, x_mask: torch.Tensor, g: Optional[torch.Tensor] = None, reverse: bool = False): def forward(
self,
x: torch.Tensor,
x_mask: torch.Tensor,
g: Optional[torch.Tensor] = None,
reverse: bool = False,
) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
x0, x1 = torch.split(x, [self.half_channels] * 2, 1) x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
h = self.pre(x0) h = self.pre(x0)
h = self.convs(h, x_mask, g=g) h = self.convs(h, x_mask, g=g)