| 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)) | |