File size: 3,887 Bytes
924c3c9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | import torch
import torch.nn as nn
class LayerNorm(nn.LayerNorm):
def __init__(self, nout, dim=-1):
super().__init__(nout, eps=1e-12)
self.dim = dim
def forward(self, x):
if self.dim == -1:
return super().forward(x)
return super().forward(x.transpose(1, -1)).transpose(1, -1)
class DurationPredictor(nn.Module):
def __init__(
self,
in_channels: int,
filter_channels: int,
n_layers: int = 2,
kernel_size: int = 3,
p_dropout: float = 0.1,
padding: str = "SAME"
):
super().__init__()
self.conv = nn.ModuleList()
self.kernel_size = kernel_size
self.padding = padding
for idx in range(n_layers):
in_chans = in_channels if idx == 0 else filter_channels
self.conv += [
nn.Sequential(
nn.ConstantPad1d(
((kernel_size - 1) // 2, (kernel_size - 1) // 2)
if padding == 'SAME' else (kernel_size - 1, 0),
0
),
nn.Conv1d(
in_chans, filter_channels,
kernel_size, stride=1, padding=0
),
nn.ReLU(),
LayerNorm(filter_channels, dim=1),
nn.Dropout(p_dropout)
)
]
self.linear = nn.Linear(filter_channels, 1)
def forward(self, x: torch.Tensor, x_mask: torch.Tensor):
x = x.transpose(1, -1)
x_mask = x_mask.unsqueeze(1).to(x.device)
for f in self.conv:
x = f(x)
x = x * x_mask.float()
x = self.linear(x.transpose(1, -1)) * x_mask.transpose(1, -1).float()
return x
class ContentAdapterBase(nn.Module):
def __init__(self, d_out):
super().__init__()
self.d_out = d_out
class CrossAttentionAdapter(ContentAdapterBase):
def __init__(
self,
d_out: int,
content_dim: int,
prefix_dim: int,
num_heads: int,
duration_predictor: DurationPredictor,
dropout: float = 0.1,
duration_grad_scale: float = 0.1,
):
super().__init__(d_out)
self.attn = nn.MultiheadAttention(
embed_dim=content_dim,
num_heads=num_heads,
dropout=dropout,
kdim=prefix_dim,
vdim=prefix_dim,
batch_first=True,
)
self.duration_grad_scale = duration_grad_scale
self.duration_predictor = duration_predictor
self.global_duration_mlp = nn.Sequential(
nn.Linear(content_dim, content_dim), nn.ReLU(),
nn.Dropout(dropout), nn.Linear(content_dim, 1)
)
self.norm = nn.LayerNorm(content_dim)
self.content_proj = nn.Conv1d(content_dim, d_out, 1)
def forward(self, content, content_mask, prefix, prefix_mask):
attn_output, attn_output_weights = self.attn(
query=content,
key=prefix,
value=prefix,
key_padding_mask=~prefix_mask.bool()
)
attn_output = attn_output * content_mask.unsqueeze(-1).float()
x = self.norm(attn_output + content)
x_grad_rescaled = x * self.duration_grad_scale + x.detach() * (
1 - self.duration_grad_scale
)
x_aggregated = (
x_grad_rescaled * content_mask.unsqueeze(-1).float()
).sum(dim=1) / content_mask.sum(dim=1, keepdim=True).float()
global_duration = self.global_duration_mlp(x_aggregated).squeeze(-1)
local_duration = self.duration_predictor(
x_grad_rescaled, content_mask
).squeeze(-1)
content = self.content_proj(x.transpose(1, 2)).transpose(1, 2)
return content, content_mask, global_duration, local_duration
|