Spaces:
Running
on
Zero
Running
on
Zero
| # MIT License | |
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # Permission is hereby granted, free of charge, to any person obtaining a copy | |
| # of this software and associated documentation files (the "Software"), to deal | |
| # in the Software without restriction, including without limitation the rights | |
| # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell | |
| # copies of the Software, and to permit persons to whom the Software is | |
| # furnished to do so, subject to the following conditions: | |
| # The above copyright notice and this permission notice shall be included in all | |
| # copies or substantial portions of the Software. | |
| # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR | |
| # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, | |
| # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE | |
| # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER | |
| # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, | |
| # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE | |
| # SOFTWARE. | |
| # Copyright (c) [2023] [Meta Platforms, Inc. and affiliates.] | |
| # Copyright (c) [2025] [Ziyue Jiang] | |
| # SPDX-License-Identifier: MIT | |
| # This file has been modified by Ziyue Jiang on 2025/03/19 | |
| # Original file was released under MIT, with the full license text # available at https://github.com/facebookresearch/encodec/blob/gh-pages/LICENSE. | |
| # This modified file is released under the same license. | |
| """Convolutional layers wrappers and utilities.""" | |
| import math | |
| import typing as tp | |
| import warnings | |
| import einops | |
| import torch | |
| from torch import nn | |
| from torch.nn import functional as F | |
| from torch.nn.utils import spectral_norm, weight_norm | |
| CONV_NORMALIZATIONS = frozenset(['none', 'weight_norm', 'spectral_norm', | |
| 'time_layer_norm', 'layer_norm', 'time_group_norm']) | |
| def apply_parametrization_norm(module: nn.Module, norm: str = 'none') -> nn.Module: | |
| assert norm in CONV_NORMALIZATIONS | |
| if norm == 'weight_norm': | |
| return weight_norm(module) | |
| elif norm == 'spectral_norm': | |
| return spectral_norm(module) | |
| else: | |
| return module | |
| def get_norm_module(module: nn.Module, causal: bool = False, norm: str = 'none', **norm_kwargs) -> nn.Module: | |
| assert norm in CONV_NORMALIZATIONS | |
| if norm == 'layer_norm': | |
| assert isinstance(module, nn.modules.conv._ConvNd) | |
| return ConvLayerNorm(module.out_channels, **norm_kwargs) | |
| elif norm == 'time_group_norm': | |
| if causal: | |
| raise ValueError("GroupNorm doesn't support causal evaluation.") | |
| assert isinstance(module, nn.modules.conv._ConvNd) | |
| return nn.GroupNorm(1, module.out_channels, **norm_kwargs) | |
| else: | |
| return nn.Identity() | |
| def get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, | |
| padding_total: int = 0) -> int: | |
| length = x.shape[-1] | |
| n_frames = (length - kernel_size + padding_total) / stride + 1 | |
| ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total) | |
| return ideal_length - length | |
| def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'zero', value: float = 0.): | |
| length = x.shape[-1] | |
| padding_left, padding_right = paddings | |
| assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right) | |
| if mode == 'reflect': | |
| max_pad = max(padding_left, padding_right) | |
| extra_pad = 0 | |
| if length <= max_pad: | |
| extra_pad = max_pad - length + 1 | |
| x = F.pad(x, (0, extra_pad)) | |
| padded = F.pad(x, paddings, mode, value) | |
| end = padded.shape[-1] - extra_pad | |
| return padded[..., :end] | |
| else: | |
| return F.pad(x, paddings, mode, value) | |
| class ConvLayerNorm(nn.LayerNorm): | |
| def __init__(self, normalized_shape: tp.Union[int, tp.List[int], torch.Size], **kwargs): | |
| super().__init__(normalized_shape, **kwargs) | |
| def forward(self, x): | |
| x = einops.rearrange(x, 'b ... t -> b t ...') | |
| x = super().forward(x) | |
| x = einops.rearrange(x, 'b t ... -> b ... t') | |
| return | |
| class NormConv1d(nn.Module): | |
| def __init__(self, *args, causal: bool = False, norm: str = 'none', | |
| norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs): | |
| super().__init__() | |
| self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm) | |
| self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs) | |
| self.norm_type = norm | |
| def forward(self, x): | |
| x = self.conv(x) | |
| x = self.norm(x) | |
| return x | |
| class SConv1d(nn.Module): | |
| def __init__(self, in_channels: int, out_channels: int, | |
| kernel_size: int, stride: int = 1, dilation: int = 1, | |
| groups: int = 1, bias: bool = True, causal: bool = False, | |
| norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {}, | |
| pad_mode: str = 'reflect'): | |
| super().__init__() | |
| # warn user on unusual setup between dilation and stride | |
| if stride > 1 and dilation > 1: | |
| warnings.warn('SConv1d has been initialized with stride > 1 and dilation > 1' | |
| f' (kernel_size={kernel_size} stride={stride}, dilation={dilation}).') | |
| self.conv = NormConv1d(in_channels, out_channels, kernel_size, stride, | |
| dilation=dilation, groups=groups, bias=bias, causal=causal, | |
| norm=norm, norm_kwargs=norm_kwargs) | |
| self.causal = causal | |
| self.pad_mode = pad_mode | |
| def forward(self, x): | |
| B, C, T = x.shape | |
| kernel_size = self.conv.conv.kernel_size[0] | |
| stride = self.conv.conv.stride[0] | |
| dilation = self.conv.conv.dilation[0] | |
| kernel_size = (kernel_size - 1) * dilation + 1 # effective kernel size with dilations | |
| padding_total = kernel_size - stride | |
| extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total) | |
| if self.causal: | |
| # Left padding for causal | |
| x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode) | |
| else: | |
| # Asymmetric padding required for odd strides | |
| padding_right = padding_total // 2 | |
| padding_left = padding_total - padding_right | |
| x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode) | |
| return self.conv(x) |