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):
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

384
style_bert_vits2/models/models.py
View File

@@ -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

369
style_bert_vits2/models/models_jp_extra.py
View File

@@ -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

42
style_bert_vits2/models/modules.py
View File

@@ -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)