File size: 1,190 Bytes
63a4356 | 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 | import torch, torch.nn as nn, torch.nn.functional as F
HOP, STRIDE = 300, 8
class StyleEncoder(nn.Module):
def __init__(self, style_dim=256, hidden=256, sr=24000):
super().__init__()
import torchaudio
self.mel = torchaudio.transforms.MelSpectrogram(sample_rate=sr, n_fft=1024, hop_length=HOP, n_mels=80, power=1.0)
self.convs = nn.ModuleList(nn.Conv1d(a, b, 5, stride=2, padding=2) for a, b in ((80, hidden), (hidden, hidden), (hidden, hidden)))
self.out = nn.Linear(hidden, style_dim)
nn.init.zeros_(self.out.weight)
nn.init.zeros_(self.out.bias)
def forward(self, audios):
device = self.out.weight.device
lengths = torch.tensor([len(a) for a in audios], device=device)
x = (torch.log(self.mel(nn.utils.rnn.pad_sequence(audios, batch_first=True).to(device)) + 1e-5) + 4) / 4
for conv in self.convs:
x = F.silu(conv(x))
frames = torch.div(lengths, HOP * STRIDE, rounding_mode="floor").clamp(min=1)
mask = (torch.arange(x.shape[-1], device=device)[None] < frames[:, None]).float()
return self.out((x * mask[:, None]).sum(-1) / mask.sum(-1, keepdim=True))
|