Spaces:
Running on Zero
Running on Zero
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the MIT license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| # Portions of this code are a modification of the fairseq project. | |
| # The original source code can be found at: | |
| # https://github.com/facebookresearch/fairseq/blob/d13e14a800bb588e5a77fb4e551f554ff9b24a72/fairseq/models/wav2vec/wav2vec.py | |
| # | |
| # Copyright (c) 2025 Yiwei Guo | |
| # Licensed under the MIT license. | |
| import math | |
| import torch.nn as nn | |
| from lscodec.models.fairseq_modules.fp32_group_norm import Fp32GroupNorm | |
| from lscodec.models.fairseq_modules.layer_norm import Fp32LayerNorm | |
| from lscodec.models.fairseq_modules.transpose_last import TransposeLast | |
| import torch | |
| def norm_block(is_layer_norm, dim, affine=True): | |
| if is_layer_norm: | |
| mod = nn.Sequential( | |
| TransposeLast(), | |
| Fp32LayerNorm(dim, elementwise_affine=affine), | |
| TransposeLast(), | |
| ) | |
| else: | |
| mod = Fp32GroupNorm(1, dim, affine=affine) | |
| return mod | |
| class ZeroPad1d(nn.Module): | |
| def __init__(self, pad_left, pad_right): | |
| super().__init__() | |
| self.pad_left = pad_left | |
| self.pad_right = pad_right | |
| def forward(self, x): | |
| return nn.functional.pad(x, (self.pad_left, self.pad_right)) | |
| class ConvFeatureExtractionModel(nn.Module): | |
| def __init__( | |
| self, | |
| conv_layers, | |
| dropout, | |
| log_compression, | |
| skip_connections, | |
| residual_scale, | |
| non_affine_group_norm, | |
| activation, | |
| ): | |
| super().__init__() | |
| def block(n_in, n_out, k, stride): | |
| return nn.Sequential( | |
| nn.Conv1d(n_in, n_out, k, stride=stride, bias=False), | |
| nn.Dropout(p=dropout), | |
| norm_block( | |
| is_layer_norm=False, dim=n_out, affine=not non_affine_group_norm | |
| ), | |
| activation, | |
| ) | |
| in_d = 1 | |
| self.conv_layers = nn.ModuleList() | |
| for dim, k, stride in conv_layers: | |
| self.conv_layers.append(block(in_d, dim, k, stride)) | |
| in_d = dim | |
| self.log_compression = log_compression | |
| self.skip_connections = skip_connections | |
| self.residual_scale = math.sqrt(residual_scale) | |
| def forward(self, x): # x is audio of shape [B, 1, T] | |
| # Bx1xT -> BxCxT | |
| for conv in self.conv_layers: | |
| residual = x | |
| x = conv(x) | |
| if self.skip_connections and x.size(1) == residual.size(1): | |
| tsz = x.size(2) | |
| r_tsz = residual.size(2) | |
| residual = residual[..., :: r_tsz // tsz][..., :tsz] | |
| x = (x + residual) * self.residual_scale | |
| if self.log_compression: | |
| x = x.abs() | |
| x = x + 1 | |
| x = x.log() | |
| return x | |
| class ConvAggregator(nn.Module): | |
| def __init__( | |
| self, | |
| conv_layers, | |
| embed, | |
| dropout, | |
| skip_connections, | |
| residual_scale, | |
| non_affine_group_norm, | |
| conv_bias, | |
| zero_pad, | |
| activation, | |
| ): | |
| super().__init__() | |
| def block(n_in, n_out, k, stride): | |
| # padding dims only really make sense for stride = 1 | |
| ka = k // 2 | |
| kb = ka - 1 if k % 2 == 0 else ka | |
| pad = ( | |
| ZeroPad1d(ka + kb, 0) if zero_pad else nn.ReplicationPad1d((ka + kb, 0)) | |
| ) | |
| return nn.Sequential( | |
| pad, | |
| nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias), | |
| nn.Dropout(p=dropout), | |
| norm_block(False, n_out, affine=not non_affine_group_norm), | |
| activation, | |
| ) | |
| in_d = embed | |
| self.conv_layers = nn.ModuleList() | |
| self.residual_proj = nn.ModuleList() | |
| for dim, k, stride in conv_layers: | |
| if in_d != dim and skip_connections: | |
| self.residual_proj.append(nn.Conv1d(in_d, dim, 1, bias=False)) | |
| else: | |
| self.residual_proj.append(None) | |
| self.conv_layers.append(block(in_d, dim, k, stride)) | |
| in_d = dim | |
| self.conv_layers = nn.Sequential(*self.conv_layers) | |
| self.skip_connections = skip_connections | |
| self.residual_scale = math.sqrt(residual_scale) | |
| def forward(self, x): | |
| for rproj, conv in zip(self.residual_proj, self.conv_layers): | |
| residual = x | |
| x = conv(x) | |
| if self.skip_connections: | |
| if rproj is not None: | |
| residual = rproj(residual) | |
| x = (x + residual) * self.residual_scale | |
| return x | |
| class LSCodecEncoder(nn.Module): | |
| def __init__(self, conv_extraction_model: ConvFeatureExtractionModel, | |
| conv_aggregation_model: ConvAggregator, | |
| dropout_p: float = 0.1): | |
| super().__init__() | |
| self.conv_models = nn.Sequential(conv_extraction_model, conv_aggregation_model) | |
| self.dropout = nn.Dropout(dropout_p) | |
| def forward(self, waveform, frame_mask, length_tolerance=3): | |
| # waveform: [B, 1, T]. Frame mask is 1 at valid values. | |
| conv_features = self.conv_models(waveform) # [B, 1024, L] | |
| conv_features = self.dropout(conv_features) | |
| B, D, L = conv_features.shape | |
| if frame_mask is None: | |
| frame_mask = torch.ones(size=(B, L)).bool().to(conv_features.device) | |
| if (frame_mask is not None) and (0 < frame_mask.shape[-1] - L <= length_tolerance): | |
| pad_tensor = torch.zeros(size=(B, D, frame_mask.shape[-1] - L)).to(conv_features.device) | |
| # NOTE: we don't pad these values with 999, but instead treat those as valid data points | |
| # NOTE: this is because the mask is longer than the valid data points, and changing the mask is inconvenient. | |
| conv_features = torch.concat([conv_features, pad_tensor], dim=-1) | |
| return conv_features | |