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

Refactor: add type hints to attentions.py / modules.py / transforms.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 22:58:40 +00:00
parent 717ba7925f
commit e1fad54c99
3 changed files with 182 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

137
style_bert_vits2/models/modules.py
View File

@@ -1,4 +1,5 @@
import math
from typing import Any, Optional, Union
import torch
from torch import nn
@@ -15,7 +16,7 @@ LRELU_SLOPE = 0.1
class LayerNorm(nn.Module):
def __init__(self, channels, eps=1e-5):
def __init__(self, channels: int, eps: float = 1e-5):
super().__init__()
self.channels = channels
self.eps = eps
@@ -23,7 +24,7 @@ class LayerNorm(nn.Module):
self.gamma = nn.Parameter(torch.ones(channels))
self.beta = nn.Parameter(torch.zeros(channels))
def forward(self, x):
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = x.transpose(1, -1)
x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
return x.transpose(1, -1)
@@ -32,12 +33,12 @@ class LayerNorm(nn.Module):
class ConvReluNorm(nn.Module):
def __init__(
self,
in_channels,
hidden_channels,
out_channels,
kernel_size,
n_layers,
p_dropout,
in_channels: int,
hidden_channels: int,
out_channels: int,
kernel_size: int,
n_layers: int,
p_dropout: float,
):
super().__init__()
self.in_channels = in_channels
@@ -69,9 +70,10 @@ class ConvReluNorm(nn.Module):
self.norm_layers.append(LayerNorm(hidden_channels))
self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
self.proj.weight.data.zero_()
assert self.proj.bias is not None
self.proj.bias.data.zero_()
def forward(self, x, x_mask):
def forward(self, x: torch.Tensor, x_mask: torch.Tensor) -> torch.Tensor:
x_org = x
for i in range(self.n_layers):
x = self.conv_layers[i](x * x_mask)
@@ -86,7 +88,7 @@ class DDSConv(nn.Module):
Dialted and Depth-Separable Convolution
"""
def __init__(self, channels, kernel_size, n_layers, p_dropout=0.0):
def __init__(self, channels: int, kernel_size: int, n_layers: int, p_dropout: float = 0.0):
super().__init__()
self.channels = channels
self.kernel_size = kernel_size
@@ -115,7 +117,7 @@ class DDSConv(nn.Module):
self.norms_1.append(LayerNorm(channels))
self.norms_2.append(LayerNorm(channels))
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:
if g is not None:
x = x + g
for i in range(self.n_layers):
@@ -133,12 +135,12 @@ class DDSConv(nn.Module):
class WN(torch.nn.Module):
def __init__(
self,
hidden_channels,
kernel_size,
dilation_rate,
n_layers,
gin_channels=0,
p_dropout=0,
hidden_channels: int,
kernel_size: int,
dilation_rate: int,
n_layers: int,
gin_channels: int = 0,
p_dropout: float = 0,
):
super(WN, self).__init__()
assert kernel_size % 2 == 1
@@ -182,7 +184,7 @@ class WN(torch.nn.Module):
res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight")
self.res_skip_layers.append(res_skip_layer)
def forward(self, x, x_mask, g=None, **kwargs):
def forward(self, x: torch.Tensor, x_mask: torch.Tensor, g: Optional[torch.Tensor] = None, **kwargs: Any) -> torch.Tensor:
output = torch.zeros_like(x)
n_channels_tensor = torch.IntTensor([self.hidden_channels])
@@ -209,7 +211,7 @@ class WN(torch.nn.Module):
output = output + res_skip_acts
return output * x_mask
def remove_weight_norm(self):
def remove_weight_norm(self) -> None:
if self.gin_channels != 0:
torch.nn.utils.remove_weight_norm(self.cond_layer)
for l in self.in_layers:
@@ -219,7 +221,7 @@ class WN(torch.nn.Module):
class ResBlock1(torch.nn.Module):
def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
def __init__(self, channels: int, kernel_size: int = 3, dilation: tuple[int, int, int] = (1, 3, 5)):
super(ResBlock1, self).__init__()
self.convs1 = nn.ModuleList(
[
@@ -293,7 +295,7 @@ class ResBlock1(torch.nn.Module):
)
self.convs2.apply(commons.init_weights)
def forward(self, x, x_mask=None):
def forward(self, x: torch.Tensor, x_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
for c1, c2 in zip(self.convs1, self.convs2):
xt = F.leaky_relu(x, LRELU_SLOPE)
if x_mask is not None:
@@ -308,7 +310,7 @@ class ResBlock1(torch.nn.Module):
x = x * x_mask
return x
def remove_weight_norm(self):
def remove_weight_norm(self) -> None:
for l in self.convs1:
remove_weight_norm(l)
for l in self.convs2:
@@ -316,7 +318,7 @@ class ResBlock1(torch.nn.Module):
class ResBlock2(torch.nn.Module):
def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
def __init__(self, channels: int, kernel_size: int = 3, dilation: tuple[int, int] = (1, 3)):
super(ResBlock2, self).__init__()
self.convs = nn.ModuleList(
[
@@ -344,7 +346,7 @@ class ResBlock2(torch.nn.Module):
)
self.convs.apply(commons.init_weights)
def forward(self, x, x_mask=None):
def forward(self, x: torch.Tensor, x_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
for c in self.convs:
xt = F.leaky_relu(x, LRELU_SLOPE)
if x_mask is not None:
@@ -355,13 +357,13 @@ class ResBlock2(torch.nn.Module):
x = x * x_mask
return x
def remove_weight_norm(self):
def remove_weight_norm(self) -> None:
for l in self.convs:
remove_weight_norm(l)
class Log(nn.Module):
def forward(self, x, x_mask, reverse=False, **kwargs):
def forward(self, x: torch.Tensor, x_mask: torch.Tensor, reverse: bool = False, **kwargs: Any):
if not reverse:
y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
logdet = torch.sum(-y, [1, 2])
@@ -372,7 +374,13 @@ class Log(nn.Module):
class Flip(nn.Module):
def forward(self, x, *args, reverse=False, **kwargs):
def forward(
self,
x: torch.Tensor,
*args: Any,
reverse: bool = False,
**kwargs: Any,
) -> Union[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
x = torch.flip(x, [1])
if not reverse:
logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
@@ -382,13 +390,19 @@ class Flip(nn.Module):
class ElementwiseAffine(nn.Module):
def __init__(self, channels):
def __init__(self, channels: int):
super().__init__()
self.channels = channels
self.m = nn.Parameter(torch.zeros(channels, 1))
self.logs = nn.Parameter(torch.zeros(channels, 1))
def forward(self, x, x_mask, reverse=False, **kwargs):
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 = self.m + torch.exp(self.logs) * x
y = y * x_mask
@@ -402,14 +416,14 @@ class ElementwiseAffine(nn.Module):
class ResidualCouplingLayer(nn.Module):
def __init__(
self,
channels,
hidden_channels,
kernel_size,
dilation_rate,
n_layers,
p_dropout=0,
gin_channels=0,
mean_only=False,
channels: int,
hidden_channels: int,
kernel_size: int,
dilation_rate: int,
n_layers: int,
p_dropout: float = 0,
gin_channels: int = 0,
mean_only: bool = False,
):
assert channels % 2 == 0, "channels should be divisible by 2"
super().__init__()
@@ -432,9 +446,10 @@ class ResidualCouplingLayer(nn.Module):
)
self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
self.post.weight.data.zero_()
assert self.post.bias is not None
self.post.bias.data.zero_()
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):
x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
h = self.pre(x0) * x_mask
h = self.enc(h, x_mask, g=g)
@@ -459,12 +474,12 @@ class ResidualCouplingLayer(nn.Module):
class ConvFlow(nn.Module):
def __init__(
self,
in_channels,
filter_channels,
kernel_size,
n_layers,
num_bins=10,
tail_bound=5.0,
in_channels: int,
filter_channels: int,
kernel_size: int,
n_layers: int,
num_bins: int = 10,
tail_bound: float = 5.0,
):
super().__init__()
self.in_channels = in_channels
@@ -481,9 +496,10 @@ class ConvFlow(nn.Module):
filter_channels, self.half_channels * (num_bins * 3 - 1), 1
)
self.proj.weight.data.zero_()
assert self.proj.bias is not None
self.proj.bias.data.zero_()
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):
x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
h = self.pre(x0)
h = self.convs(h, x_mask, g=g)
@@ -519,17 +535,17 @@ class ConvFlow(nn.Module):
class TransformerCouplingLayer(nn.Module):
def __init__(
self,
channels,
hidden_channels,
kernel_size,
n_layers,
n_heads,
p_dropout=0,
filter_channels=0,
mean_only=False,
wn_sharing_parameter=None,
gin_channels=0,
):
channels: int,
hidden_channels: int,
kernel_size: int,
n_layers: int,
n_heads: int,
p_dropout: float = 0,
filter_channels: int = 0,
mean_only: bool = False,
wn_sharing_parameter: Optional[nn.Module] = None,
gin_channels: int = 0,
) -> None:
assert channels % 2 == 0, "channels should be divisible by 2"
super().__init__()
self.channels = channels
@@ -556,9 +572,16 @@ class TransformerCouplingLayer(nn.Module):
)
self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
self.post.weight.data.zero_()
assert self.post.bias is not None
self.post.bias.data.zero_()
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,
) -> 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)