File size: 2,074 Bytes
36a67ca |
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 |
import torch
import torch.nn as nn
from .inference_hubert import InferenceHubertBase
from .vae_memory_bank import VAEMemoryBank
def create_padding_mask(waveforms_lengths: torch.Tensor = None):
if waveforms_lengths is None:
return None
batch = waveforms_lengths.shape[0]
max_len = waveforms_lengths.max()
device = waveforms_lengths.device
padding_mask = torch.ones(batch, max_len, dtype=torch.bool, device=device)
for idx, length in enumerate(waveforms_lengths):
padding_mask[idx, :length] = 0
return padding_mask
def unfreeze_layers(model: nn.Module, root_name: str):
for name, param in model.named_parameters():
if root_name in name[: len(root_name)]:
param.requires_grad = True
class PosteriorHubert(nn.Module):
def __init__(
self, out_channels, feature_channels, downsample_channels, output_layer=11
) -> None:
super().__init__()
self.out_channels = out_channels
self.feature_channels = feature_channels
self.downsample_channels = downsample_channels
self.output_layer = output_layer
self.hubert = InferenceHubertBase()
self.memory_bank = VAEMemoryBank(
n_hidden_dims=feature_channels,
bank_size=1000,
output_channels=downsample_channels,
)
self.proj = nn.Conv1d(downsample_channels, out_channels * 2, 1)
def forward(self, waveforms: torch.Tensor, waveforms_lengths: torch.Tensor, g=None):
features, features_mask = self.hubert.extract_features(
source=waveforms,
padding_mask=create_padding_mask(waveforms_lengths),
output_layer=self.output_layer,
)
x = self.memory_bank(features.transpose(1, 2))
x_mask = (~features_mask).unsqueeze(1).to(torch.float32)
x = x[:, :, : x_mask.shape[-1]]
stats = self.proj(x) * x_mask
m, logs = torch.split(stats, self.out_channels, dim=1)
z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask
return z, m, logs, x_mask
|