diff --git a/style_bert_vits2/models/attentions.py b/style_bert_vits2/models/attentions.py index b262681..03b238d 100644 --- a/style_bert_vits2/models/attentions.py +++ b/style_bert_vits2/models/attentions.py @@ -9,7 +9,7 @@ from style_bert_vits2.models import commons 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__() self.channels = channels self.eps = eps @@ -45,7 +45,7 @@ class Encoder(nn.Module): window_size: int = 4, isflow: bool = True, **kwargs: Any - ): + ) -> None: super().__init__() self.hidden_channels = hidden_channels self.filter_channels = filter_channels @@ -132,7 +132,7 @@ class Decoder(nn.Module): proximal_bias: bool = False, proximal_init: bool = True, **kwargs: Any - ): + ) -> None: super().__init__() self.hidden_channels = hidden_channels self.filter_channels = filter_channels @@ -180,7 +180,7 @@ class Decoder(nn.Module): ) 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 h: encoder output @@ -218,7 +218,7 @@ class MultiHeadAttention(nn.Module): block_length: Optional[int] = None, proximal_bias: bool = False, proximal_init: bool = False, - ): + ) -> None: super().__init__() assert channels % n_heads == 0 @@ -272,7 +272,13 @@ class MultiHeadAttention(nn.Module): x = self.conv_o(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] 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) @@ -419,7 +425,7 @@ class FFN(nn.Module): p_dropout: float = 0.0, activation: Optional[str] = None, causal: bool = False, - ): + ) -> None: super().__init__() self.in_channels = in_channels self.out_channels = out_channels diff --git a/style_bert_vits2/models/models.py b/style_bert_vits2/models/models.py index 4efd100..829ca4a 100644 --- a/style_bert_vits2/models/models.py +++ b/style_bert_vits2/models/models.py @@ -1,4 +1,5 @@ import math +from typing import Any, Optional import torch 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 modules 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 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 @@ -51,7 +56,13 @@ class DurationDiscriminator(nn.Module): # vits2 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) x = torch.cat([x, dur], dim=1) x = self.pre_out_conv_1(x * x_mask) @@ -67,7 +78,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) @@ -92,17 +110,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 @@ -114,16 +132,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 ) @@ -145,7 +164,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) @@ -158,13 +183,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 @@ -204,7 +229,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: @@ -268,8 +301,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 @@ -292,7 +330,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) @@ -312,17 +350,17 @@ class DurationPredictor(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, - n_speakers, - 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, + n_speakers: int, + gin_channels: int = 0, + ) -> None: super().__init__() self.n_vocab = n_vocab self.out_channels = out_channels @@ -357,17 +395,17 @@ class TextEncoder(nn.Module): def forward( self, - x, - x_lengths, - tone, - language, - bert, - ja_bert, - en_bert, - style_vec, - sid, - g=None, - ): + x: torch.Tensor, + x_lengths: torch.Tensor, + tone: torch.Tensor, + language: torch.Tensor, + bert: torch.Tensor, + ja_bert: torch.Tensor, + en_bert: torch.Tensor, + style_vec: torch.Tensor, + sid: 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) ja_bert_emb = self.ja_bert_proj(ja_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): 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 @@ -431,7 +469,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) @@ -444,14 +488,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 @@ -471,7 +515,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 ) @@ -486,22 +535,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)): @@ -518,20 +567,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(init_weights) + 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) @@ -545,6 +596,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) @@ -552,7 +604,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) @@ -561,7 +613,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 @@ -574,7 +626,7 @@ class DiscriminatorP(torch.nn.Module): 32, (kernel_size, 1), (stride, 1), - padding=(get_padding(kernel_size, 1), 0), + padding=(commons.get_padding(kernel_size, 1), 0), ) ), norm_f( @@ -583,7 +635,7 @@ class DiscriminatorP(torch.nn.Module): 128, (kernel_size, 1), (stride, 1), - padding=(get_padding(kernel_size, 1), 0), + padding=(commons.get_padding(kernel_size, 1), 0), ) ), norm_f( @@ -592,7 +644,7 @@ class DiscriminatorP(torch.nn.Module): 512, (kernel_size, 1), (stride, 1), - padding=(get_padding(kernel_size, 1), 0), + padding=(commons.get_padding(kernel_size, 1), 0), ) ), norm_f( @@ -601,7 +653,7 @@ class DiscriminatorP(torch.nn.Module): 1024, (kernel_size, 1), (stride, 1), - padding=(get_padding(kernel_size, 1), 0), + padding=(commons.get_padding(kernel_size, 1), 0), ) ), norm_f( @@ -610,14 +662,14 @@ class DiscriminatorP(torch.nn.Module): 1024, (kernel_size, 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))) - def forward(self, x): + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: fmap = [] # 1d to 2d @@ -640,7 +692,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( @@ -655,7 +707,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: @@ -670,7 +722,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] @@ -680,7 +732,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 = [] @@ -702,7 +758,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] @@ -731,7 +787,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: @@ -749,7 +805,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 @@ -762,31 +818,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=4, - 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 = 4, + 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 @@ -884,18 +940,27 @@ class SynthesizerTrn(nn.Module): def forward( self, - x, - x_lengths, - y, - y_lengths, - sid, - tone, - language, - bert, - ja_bert, - en_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, + ja_bert: torch.Tensor, + en_bert: torch.Tensor, + 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: g = self.emb_g(sid).unsqueeze(-1) # [b, h, 1] else: @@ -973,27 +1038,28 @@ class SynthesizerTrn(nn.Module): def infer( self, - x, - x_lengths, - sid, - tone, - language, - bert, - ja_bert, - en_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, + ja_bert: torch.Tensor, + en_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, ja_bert, en_bert, style_vec, sid, g=g diff --git a/style_bert_vits2/models/models_jp_extra.py b/style_bert_vits2/models/models_jp_extra.py index 3a43d51..a43a715 100644 --- a/style_bert_vits2/models/models_jp_extra.py +++ b/style_bert_vits2/models/models_jp_extra.py @@ -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 diff --git a/style_bert_vits2/models/modules.py b/style_bert_vits2/models/modules.py index df38071..8eed963 100644 --- a/style_bert_vits2/models/modules.py +++ b/style_bert_vits2/models/modules.py @@ -16,7 +16,7 @@ LRELU_SLOPE = 0.1 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__() self.channels = channels self.eps = eps @@ -39,7 +39,7 @@ class ConvReluNorm(nn.Module): kernel_size: int, n_layers: int, p_dropout: float, - ): + ) -> None: super().__init__() self.in_channels = in_channels self.hidden_channels = hidden_channels @@ -88,7 +88,7 @@ class DDSConv(nn.Module): 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__() self.channels = channels self.kernel_size = kernel_size @@ -141,7 +141,7 @@ class WN(torch.nn.Module): n_layers: int, gin_channels: int = 0, p_dropout: float = 0, - ): + ) -> None: super(WN, self).__init__() assert kernel_size % 2 == 1 self.hidden_channels = hidden_channels @@ -221,7 +221,7 @@ class WN(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__() self.convs1 = nn.ModuleList( [ @@ -318,7 +318,7 @@ class ResBlock1(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__() self.convs = nn.ModuleList( [ @@ -363,7 +363,13 @@ class ResBlock2(torch.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: y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask logdet = torch.sum(-y, [1, 2]) @@ -390,7 +396,7 @@ class Flip(nn.Module): class ElementwiseAffine(nn.Module): - def __init__(self, channels: int): + def __init__(self, channels: int) -> None: super().__init__() self.channels = channels self.m = nn.Parameter(torch.zeros(channels, 1)) @@ -424,7 +430,7 @@ class ResidualCouplingLayer(nn.Module): p_dropout: float = 0, gin_channels: int = 0, mean_only: bool = False, - ): + ) -> None: assert channels % 2 == 0, "channels should be divisible by 2" super().__init__() self.channels = channels @@ -449,7 +455,13 @@ class ResidualCouplingLayer(nn.Module): assert self.post.bias is not None 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) h = self.pre(x0) * x_mask h = self.enc(h, x_mask, g=g) @@ -480,7 +492,7 @@ class ConvFlow(nn.Module): n_layers: int, num_bins: int = 10, tail_bound: float = 5.0, - ): + ) -> None: super().__init__() self.in_channels = in_channels self.filter_channels = filter_channels @@ -499,7 +511,13 @@ class ConvFlow(nn.Module): assert self.proj.bias is not None 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) h = self.pre(x0) h = self.convs(h, x_mask, g=g)