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