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 torch
from torch import nn
@@ -7,7 +9,7 @@ from style_bert_vits2.models import commons
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
@@ -15,14 +17,14 @@ 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)
@torch.jit.script
def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
@torch.jit.script # type: ignore
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]
in_act = input_a + input_b
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):
def __init__(
self,
hidden_channels,
filter_channels,
n_heads,
n_layers,
kernel_size=1,
p_dropout=0.0,
window_size=4,
isflow=True,
**kwargs
hidden_channels: int,
filter_channels: int,
n_heads: int,
n_layers: int,
kernel_size: int = 1,
p_dropout: float = 0.0,
window_size: int = 4,
isflow: bool = True,
**kwargs: Any
):
super().__init__()
self.hidden_channels = hidden_channels
@@ -97,12 +99,13 @@ class Encoder(nn.Module):
)
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)
x = x * x_mask
for i in range(self.n_layers):
if i == self.cond_layer_idx and g is not None:
g = self.spk_emb_linear(g.transpose(1, 2))
assert g is not None
g = g.transpose(1, 2)
x = x + g
x = x * x_mask
@@ -120,15 +123,15 @@ class Encoder(nn.Module):
class Decoder(nn.Module):
def __init__(
self,
hidden_channels,
filter_channels,
n_heads,
n_layers,
kernel_size=1,
p_dropout=0.0,
proximal_bias=False,
proximal_init=True,
**kwargs
hidden_channels: int,
filter_channels: int,
n_heads: int,
n_layers: int,
kernel_size: int = 1,
p_dropout: float = 0.0,
proximal_bias: bool = False,
proximal_init: bool = True,
**kwargs: Any
):
super().__init__()
self.hidden_channels = hidden_channels
@@ -177,7 +180,7 @@ class Decoder(nn.Module):
)
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
h: encoder output
@@ -206,15 +209,15 @@ class Decoder(nn.Module):
class MultiHeadAttention(nn.Module):
def __init__(
self,
channels,
out_channels,
n_heads,
p_dropout=0.0,
window_size=None,
heads_share=True,
block_length=None,
proximal_bias=False,
proximal_init=False,
channels: int,
out_channels: int,
n_heads: int,
p_dropout: float = 0.0,
window_size: Optional[int] = None,
heads_share: bool = True,
block_length: Optional[int] = None,
proximal_bias: bool = False,
proximal_init: bool = False,
):
super().__init__()
assert channels % n_heads == 0
@@ -255,9 +258,11 @@ class MultiHeadAttention(nn.Module):
if proximal_init:
with torch.no_grad():
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)
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)
k = self.conv_k(c)
v = self.conv_v(c)
@@ -267,7 +272,7 @@ class MultiHeadAttention(nn.Module):
x = self.conv_o(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]
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)
@@ -318,7 +323,7 @@ class MultiHeadAttention(nn.Module):
) # [b, n_h, t_t, d_k] -> [b, d, t_t]
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]
y: [h or 1, m, d]
@@ -327,7 +332,7 @@ class MultiHeadAttention(nn.Module):
ret = torch.matmul(x, y.unsqueeze(0))
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]
y: [h or 1, m, d]
@@ -336,8 +341,9 @@ class MultiHeadAttention(nn.Module):
ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
return ret
def _get_relative_embeddings(self, relative_embeddings, length):
2 * self.window_size + 1
def _get_relative_embeddings(self, relative_embeddings: torch.Tensor, length: int) -> torch.Tensor:
assert self.window_size is not None
2 * self.window_size + 1 # type: ignore
# Pad first before slice to avoid using cond ops.
pad_length = max(length - (self.window_size + 1), 0)
slice_start_position = max((self.window_size + 1) - length, 0)
@@ -354,7 +360,7 @@ class MultiHeadAttention(nn.Module):
]
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]
ret: [b, h, l, l]
@@ -375,7 +381,7 @@ class MultiHeadAttention(nn.Module):
]
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]
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:]
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.
Args:
length: an integer scalar.
@@ -406,13 +412,13 @@ class MultiHeadAttention(nn.Module):
class FFN(nn.Module):
def __init__(
self,
in_channels,
out_channels,
filter_channels,
kernel_size,
p_dropout=0.0,
activation=None,
causal=False,
in_channels: int,
out_channels: int,
filter_channels: int,
kernel_size: int,
p_dropout: float = 0.0,
activation: Optional[str] = None,
causal: bool = False,
):
super().__init__()
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.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))
if self.activation == "gelu":
x = x * torch.sigmoid(1.702 * x)
@@ -442,7 +448,7 @@ class FFN(nn.Module):
x = self.conv_2(self.padding(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:
return x
pad_l = self.kernel_size - 1
@@ -451,7 +457,7 @@ class FFN(nn.Module):
x = F.pad(x, commons.convert_pad_shape(padding))
return x
def _same_padding(self, x):
def _same_padding(self, x: torch.Tensor) -> torch.Tensor:
if self.kernel_size == 1:
return x
pad_l = (self.kernel_size - 1) // 2