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

110
style_bert_vits2/models/attentions.py
View File

@@ -1,3 +1,5 @@
from typing import Any, Optional
import math import math
import torch import torch
from torch import nn from torch import nn
@@ -7,7 +9,7 @@ from style_bert_vits2.models import commons
class LayerNorm(nn.Module): class LayerNorm(nn.Module):
def __init__(self, channels, eps=1e-5): def __init__(self, channels: int, eps: float = 1e-5):
super().__init__() super().__init__()
self.channels = channels self.channels = channels
self.eps = eps self.eps = eps
@@ -15,14 +17,14 @@ class LayerNorm(nn.Module):
self.gamma = nn.Parameter(torch.ones(channels)) self.gamma = nn.Parameter(torch.ones(channels))
self.beta = nn.Parameter(torch.zeros(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 = x.transpose(1, -1)
x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps) x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
return x.transpose(1, -1) return x.transpose(1, -1)
@torch.jit.script @torch.jit.script # type: ignore
def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels): def fused_add_tanh_sigmoid_multiply(input_a: torch.Tensor, input_b: torch.Tensor, n_channels: list[int]) -> torch.Tensor:
n_channels_int = n_channels[0] n_channels_int = n_channels[0]
in_act = input_a + input_b in_act = input_a + input_b
t_act = torch.tanh(in_act[:, :n_channels_int, :]) t_act = torch.tanh(in_act[:, :n_channels_int, :])
@@ -34,15 +36,15 @@ def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
class Encoder(nn.Module): class Encoder(nn.Module):
def __init__( def __init__(
self, self,
hidden_channels, hidden_channels: int,
filter_channels, filter_channels: int,
n_heads, n_heads: int,
n_layers, n_layers: int,
kernel_size=1, kernel_size: int = 1,
p_dropout=0.0, p_dropout: float = 0.0,
window_size=4, window_size: int = 4,
isflow=True, isflow: bool = True,
**kwargs **kwargs: Any
): ):
super().__init__() super().__init__()
self.hidden_channels = hidden_channels self.hidden_channels = hidden_channels
@@ -97,12 +99,13 @@ class Encoder(nn.Module):
) )
self.norm_layers_2.append(LayerNorm(hidden_channels)) self.norm_layers_2.append(LayerNorm(hidden_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:
attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1) attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
x = x * x_mask x = x * x_mask
for i in range(self.n_layers): for i in range(self.n_layers):
if i == self.cond_layer_idx and g is not None: if i == self.cond_layer_idx and g is not None:
g = self.spk_emb_linear(g.transpose(1, 2)) g = self.spk_emb_linear(g.transpose(1, 2))
assert g is not None
g = g.transpose(1, 2) g = g.transpose(1, 2)
x = x + g x = x + g
x = x * x_mask x = x * x_mask
@@ -120,15 +123,15 @@ class Encoder(nn.Module):
class Decoder(nn.Module): class Decoder(nn.Module):
def __init__( def __init__(
self, self,
hidden_channels, hidden_channels: int,
filter_channels, filter_channels: int,
n_heads, n_heads: int,
n_layers, n_layers: int,
kernel_size=1, kernel_size: int = 1,
p_dropout=0.0, p_dropout: float = 0.0,
proximal_bias=False, proximal_bias: bool = False,
proximal_init=True, proximal_init: bool = True,
**kwargs **kwargs: Any
): ):
super().__init__() super().__init__()
self.hidden_channels = hidden_channels self.hidden_channels = hidden_channels
@@ -177,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, x_mask, h, h_mask): def forward(self, x: torch.Tensor, x_mask: torch.Tensor, h: torch.Tensor, h_mask: torch.Tensor):
""" """
x: decoder input x: decoder input
h: encoder output h: encoder output
@@ -206,15 +209,15 @@ class Decoder(nn.Module):
class MultiHeadAttention(nn.Module): class MultiHeadAttention(nn.Module):
def __init__( def __init__(
self, self,
channels, channels: int,
out_channels, out_channels: int,
n_heads, n_heads: int,
p_dropout=0.0, p_dropout: float = 0.0,
window_size=None, window_size: Optional[int] = None,
heads_share=True, heads_share: bool = True,
block_length=None, block_length: Optional[int] = None,
proximal_bias=False, proximal_bias: bool = False,
proximal_init=False, proximal_init: bool = False,
): ):
super().__init__() super().__init__()
assert channels % n_heads == 0 assert channels % n_heads == 0
@@ -255,9 +258,11 @@ class MultiHeadAttention(nn.Module):
if proximal_init: if proximal_init:
with torch.no_grad(): with torch.no_grad():
self.conv_k.weight.copy_(self.conv_q.weight) self.conv_k.weight.copy_(self.conv_q.weight)
assert self.conv_k.bias is not None
assert self.conv_q.bias is not None
self.conv_k.bias.copy_(self.conv_q.bias) self.conv_k.bias.copy_(self.conv_q.bias)
def forward(self, x, c, attn_mask=None): def forward(self, x: torch.Tensor, c: torch.Tensor, attn_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
q = self.conv_q(x) q = self.conv_q(x)
k = self.conv_k(c) k = self.conv_k(c)
v = self.conv_v(c) v = self.conv_v(c)
@@ -267,7 +272,7 @@ class MultiHeadAttention(nn.Module):
x = self.conv_o(x) x = self.conv_o(x)
return x return x
def attention(self, query, key, value, mask=None): 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)
@@ -318,7 +323,7 @@ class MultiHeadAttention(nn.Module):
) # [b, n_h, t_t, d_k] -> [b, d, t_t] ) # [b, n_h, t_t, d_k] -> [b, d, t_t]
return output, p_attn return output, p_attn
def _matmul_with_relative_values(self, x, y): def _matmul_with_relative_values(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
""" """
x: [b, h, l, m] x: [b, h, l, m]
y: [h or 1, m, d] y: [h or 1, m, d]
@@ -327,7 +332,7 @@ class MultiHeadAttention(nn.Module):
ret = torch.matmul(x, y.unsqueeze(0)) ret = torch.matmul(x, y.unsqueeze(0))
return ret return ret
def _matmul_with_relative_keys(self, x, y): def _matmul_with_relative_keys(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
""" """
x: [b, h, l, d] x: [b, h, l, d]
y: [h or 1, m, d] y: [h or 1, m, d]
@@ -336,8 +341,9 @@ class MultiHeadAttention(nn.Module):
ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1)) ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
return ret return ret
def _get_relative_embeddings(self, relative_embeddings, length): def _get_relative_embeddings(self, relative_embeddings: torch.Tensor, length: int) -> torch.Tensor:
2 * self.window_size + 1 assert self.window_size is not None
2 * self.window_size + 1 # type: ignore
# Pad first before slice to avoid using cond ops. # Pad first before slice to avoid using cond ops.
pad_length = max(length - (self.window_size + 1), 0) pad_length = max(length - (self.window_size + 1), 0)
slice_start_position = max((self.window_size + 1) - length, 0) slice_start_position = max((self.window_size + 1) - length, 0)
@@ -354,7 +360,7 @@ class MultiHeadAttention(nn.Module):
] ]
return used_relative_embeddings return used_relative_embeddings
def _relative_position_to_absolute_position(self, x): def _relative_position_to_absolute_position(self, x: torch.Tensor) -> torch.Tensor:
""" """
x: [b, h, l, 2*l-1] x: [b, h, l, 2*l-1]
ret: [b, h, l, l] ret: [b, h, l, l]
@@ -375,7 +381,7 @@ class MultiHeadAttention(nn.Module):
] ]
return x_final return x_final
def _absolute_position_to_relative_position(self, x): def _absolute_position_to_relative_position(self, x: torch.Tensor) -> torch.Tensor:
""" """
x: [b, h, l, l] x: [b, h, l, l]
ret: [b, h, l, 2*l-1] ret: [b, h, l, 2*l-1]
@@ -391,7 +397,7 @@ class MultiHeadAttention(nn.Module):
x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:] x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:]
return x_final return x_final
def _attention_bias_proximal(self, length): def _attention_bias_proximal(self, length: int) -> torch.Tensor:
"""Bias for self-attention to encourage attention to close positions. """Bias for self-attention to encourage attention to close positions.
Args: Args:
length: an integer scalar. length: an integer scalar.
@@ -406,13 +412,13 @@ class MultiHeadAttention(nn.Module):
class FFN(nn.Module): class FFN(nn.Module):
def __init__( def __init__(
self, self,
in_channels, in_channels: int,
out_channels, out_channels: int,
filter_channels, filter_channels: int,
kernel_size, kernel_size: int,
p_dropout=0.0, p_dropout: float = 0.0,
activation=None, activation: Optional[str] = None,
causal=False, causal: bool = False,
): ):
super().__init__() super().__init__()
self.in_channels = in_channels self.in_channels = in_channels
@@ -432,7 +438,7 @@ class FFN(nn.Module):
self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size) self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
self.drop = nn.Dropout(p_dropout) self.drop = nn.Dropout(p_dropout)
def forward(self, x, x_mask): def forward(self, x: torch.Tensor, x_mask: torch.Tensor) -> torch.Tensor:
x = self.conv_1(self.padding(x * x_mask)) x = self.conv_1(self.padding(x * x_mask))
if self.activation == "gelu": if self.activation == "gelu":
x = x * torch.sigmoid(1.702 * x) x = x * torch.sigmoid(1.702 * x)
@@ -442,7 +448,7 @@ class FFN(nn.Module):
x = self.conv_2(self.padding(x * x_mask)) x = self.conv_2(self.padding(x * x_mask))
return x * x_mask return x * x_mask
def _causal_padding(self, x): def _causal_padding(self, x: torch.Tensor) -> torch.Tensor:
if self.kernel_size == 1: if self.kernel_size == 1:
return x return x
pad_l = self.kernel_size - 1 pad_l = self.kernel_size - 1
@@ -451,7 +457,7 @@ class FFN(nn.Module):
x = F.pad(x, commons.convert_pad_shape(padding)) x = F.pad(x, commons.convert_pad_shape(padding))
return x return x
def _same_padding(self, x): def _same_padding(self, x: torch.Tensor) -> torch.Tensor:
if self.kernel_size == 1: if self.kernel_size == 1:
return x return x
pad_l = (self.kernel_size - 1) // 2 pad_l = (self.kernel_size - 1) // 2

137
style_bert_vits2/models/modules.py
View File

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

84
style_bert_vits2/models/transforms.py
View File

@@ -1,7 +1,8 @@
import torch from typing import Optional
from torch.nn import functional as F
import numpy as np import numpy as np
import torch
from torch.nn import functional as F
DEFAULT_MIN_BIN_WIDTH = 1e-3 DEFAULT_MIN_BIN_WIDTH = 1e-3
@@ -10,17 +11,18 @@ DEFAULT_MIN_DERIVATIVE = 1e-3
def piecewise_rational_quadratic_transform( def piecewise_rational_quadratic_transform(
inputs, inputs: torch.Tensor,
unnormalized_widths, unnormalized_widths: torch.Tensor,
unnormalized_heights, unnormalized_heights: torch.Tensor,
unnormalized_derivatives, unnormalized_derivatives: torch.Tensor,
inverse=False, inverse: bool = False,
tails=None, tails: Optional[str] = None,
tail_bound=1.0, tail_bound: float = 1.0,
min_bin_width=DEFAULT_MIN_BIN_WIDTH, min_bin_width: float = DEFAULT_MIN_BIN_WIDTH,
min_bin_height=DEFAULT_MIN_BIN_HEIGHT, min_bin_height: float = DEFAULT_MIN_BIN_HEIGHT,
min_derivative=DEFAULT_MIN_DERIVATIVE, min_derivative: float = DEFAULT_MIN_DERIVATIVE,
): ) -> tuple[torch.Tensor, torch.Tensor]:
if tails is None: if tails is None:
spline_fn = rational_quadratic_spline spline_fn = rational_quadratic_spline
spline_kwargs = {} spline_kwargs = {}
@@ -37,28 +39,29 @@ def piecewise_rational_quadratic_transform(
min_bin_width=min_bin_width, min_bin_width=min_bin_width,
min_bin_height=min_bin_height, min_bin_height=min_bin_height,
min_derivative=min_derivative, min_derivative=min_derivative,
**spline_kwargs **spline_kwargs # type: ignore
) )
return outputs, logabsdet return outputs, logabsdet
def searchsorted(bin_locations, inputs, eps=1e-6): def searchsorted(bin_locations: torch.Tensor, inputs: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
bin_locations[..., -1] += eps bin_locations[..., -1] += eps
return torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1 return torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1
def unconstrained_rational_quadratic_spline( def unconstrained_rational_quadratic_spline(
inputs, inputs: torch.Tensor,
unnormalized_widths, unnormalized_widths: torch.Tensor,
unnormalized_heights, unnormalized_heights: torch.Tensor,
unnormalized_derivatives, unnormalized_derivatives: torch.Tensor,
inverse=False, inverse: bool = False,
tails="linear", tails: str = "linear",
tail_bound=1.0, tail_bound: float = 1.0,
min_bin_width=DEFAULT_MIN_BIN_WIDTH, min_bin_width: float = DEFAULT_MIN_BIN_WIDTH,
min_bin_height=DEFAULT_MIN_BIN_HEIGHT, min_bin_height: float = DEFAULT_MIN_BIN_HEIGHT,
min_derivative=DEFAULT_MIN_DERIVATIVE, min_derivative: float = DEFAULT_MIN_DERIVATIVE,
): ) -> tuple[torch.Tensor, torch.Tensor]:
inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound) inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
outside_interval_mask = ~inside_interval_mask outside_interval_mask = ~inside_interval_mask
@@ -74,7 +77,7 @@ def unconstrained_rational_quadratic_spline(
outputs[outside_interval_mask] = inputs[outside_interval_mask] outputs[outside_interval_mask] = inputs[outside_interval_mask]
logabsdet[outside_interval_mask] = 0 logabsdet[outside_interval_mask] = 0
else: else:
raise RuntimeError("{} tails are not implemented.".format(tails)) raise RuntimeError(f"{tails} tails are not implemented.")
( (
outputs[inside_interval_mask], outputs[inside_interval_mask],
@@ -98,19 +101,20 @@ def unconstrained_rational_quadratic_spline(
def rational_quadratic_spline( def rational_quadratic_spline(
inputs, inputs: torch.Tensor,
unnormalized_widths, unnormalized_widths: torch.Tensor,
unnormalized_heights, unnormalized_heights: torch.Tensor,
unnormalized_derivatives, unnormalized_derivatives: torch.Tensor,
inverse=False, inverse: bool = False,
left=0.0, left: float = 0.0,
right=1.0, right: float = 1.0,
bottom=0.0, bottom: float = 0.0,
top=1.0, top: float = 1.0,
min_bin_width=DEFAULT_MIN_BIN_WIDTH, min_bin_width: float = DEFAULT_MIN_BIN_WIDTH,
min_bin_height=DEFAULT_MIN_BIN_HEIGHT, min_bin_height: float = DEFAULT_MIN_BIN_HEIGHT,
min_derivative=DEFAULT_MIN_DERIVATIVE, min_derivative: float = DEFAULT_MIN_DERIVATIVE,
): ) -> tuple[torch.Tensor, torch.Tensor]:
if torch.min(inputs) < left or torch.max(inputs) > right: if torch.min(inputs) < left or torch.max(inputs) > right:
raise ValueError("Input to a transform is not within its domain") raise ValueError("Input to a transform is not within its domain")