refactor: rename canto-backend → backend, canto-frontend → frontend
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
292
backend/indextts/s2mel/modules/encodec.py
Normal file
292
backend/indextts/s2mel/modules/encodec.py
Normal file
@@ -0,0 +1,292 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
"""Convolutional layers wrappers and utilities."""
|
||||
|
||||
import math
|
||||
import typing as tp
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
from torch.nn.utils import spectral_norm, weight_norm
|
||||
|
||||
import typing as tp
|
||||
|
||||
import einops
|
||||
|
||||
|
||||
class ConvLayerNorm(nn.LayerNorm):
|
||||
"""
|
||||
Convolution-friendly LayerNorm that moves channels to last dimensions
|
||||
before running the normalization and moves them back to original position right after.
|
||||
"""
|
||||
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
|
||||
|
||||
|
||||
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:
|
||||
# We already check was in CONV_NORMALIZATION, so any other choice
|
||||
# doesn't need reparametrization.
|
||||
return module
|
||||
|
||||
|
||||
def get_norm_module(module: nn.Module, causal: bool = False, norm: str = 'none', **norm_kwargs) -> nn.Module:
|
||||
"""Return the proper normalization module. If causal is True, this will ensure the returned
|
||||
module is causal, or return an error if the normalization doesn't support causal evaluation.
|
||||
"""
|
||||
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:
|
||||
"""See `pad_for_conv1d`.
|
||||
"""
|
||||
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 pad_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0):
|
||||
"""Pad for a convolution to make sure that the last window is full.
|
||||
Extra padding is added at the end. This is required to ensure that we can rebuild
|
||||
an output of the same length, as otherwise, even with padding, some time steps
|
||||
might get removed.
|
||||
For instance, with total padding = 4, kernel size = 4, stride = 2:
|
||||
0 0 1 2 3 4 5 0 0 # (0s are padding)
|
||||
1 2 3 # (output frames of a convolution, last 0 is never used)
|
||||
0 0 1 2 3 4 5 0 # (output of tr. conv., but pos. 5 is going to get removed as padding)
|
||||
1 2 3 4 # once you removed padding, we are missing one time step !
|
||||
"""
|
||||
extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
|
||||
return F.pad(x, (0, extra_padding))
|
||||
|
||||
|
||||
def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'zero', value: float = 0.):
|
||||
"""Tiny wrapper around F.pad, just to allow for reflect padding on small input.
|
||||
If this is the case, we insert extra 0 padding to the right before the reflection happen.
|
||||
"""
|
||||
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)
|
||||
|
||||
|
||||
def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):
|
||||
"""Remove padding from x, handling properly zero padding. Only for 1d!"""
|
||||
padding_left, padding_right = paddings
|
||||
assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
|
||||
assert (padding_left + padding_right) <= x.shape[-1]
|
||||
end = x.shape[-1] - padding_right
|
||||
return x[..., padding_left: end]
|
||||
|
||||
|
||||
class NormConv1d(nn.Module):
|
||||
"""Wrapper around Conv1d and normalization applied to this conv
|
||||
to provide a uniform interface across normalization approaches.
|
||||
"""
|
||||
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 NormConv2d(nn.Module):
|
||||
"""Wrapper around Conv2d and normalization applied to this conv
|
||||
to provide a uniform interface across normalization approaches.
|
||||
"""
|
||||
def __init__(self, *args, norm: str = 'none',
|
||||
norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
|
||||
super().__init__()
|
||||
self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm)
|
||||
self.norm = get_norm_module(self.conv, causal=False, norm=norm, **norm_kwargs)
|
||||
self.norm_type = norm
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
x = self.norm(x)
|
||||
return x
|
||||
|
||||
|
||||
class NormConvTranspose1d(nn.Module):
|
||||
"""Wrapper around ConvTranspose1d and normalization applied to this conv
|
||||
to provide a uniform interface across normalization approaches.
|
||||
"""
|
||||
def __init__(self, *args, causal: bool = False, norm: str = 'none',
|
||||
norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
|
||||
super().__init__()
|
||||
self.convtr = apply_parametrization_norm(nn.ConvTranspose1d(*args, **kwargs), norm)
|
||||
self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs)
|
||||
self.norm_type = norm
|
||||
|
||||
def forward(self, x):
|
||||
x = self.convtr(x)
|
||||
x = self.norm(x)
|
||||
return x
|
||||
|
||||
|
||||
class NormConvTranspose2d(nn.Module):
|
||||
"""Wrapper around ConvTranspose2d and normalization applied to this conv
|
||||
to provide a uniform interface across normalization approaches.
|
||||
"""
|
||||
def __init__(self, *args, norm: str = 'none',
|
||||
norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
|
||||
super().__init__()
|
||||
self.convtr = apply_parametrization_norm(nn.ConvTranspose2d(*args, **kwargs), norm)
|
||||
self.norm = get_norm_module(self.convtr, causal=False, norm=norm, **norm_kwargs)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.convtr(x)
|
||||
x = self.norm(x)
|
||||
return x
|
||||
|
||||
|
||||
class SConv1d(nn.Module):
|
||||
"""Conv1d with some builtin handling of asymmetric or causal padding
|
||||
and normalization.
|
||||
"""
|
||||
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', **kwargs):
|
||||
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)
|
||||
|
||||
|
||||
class SConvTranspose1d(nn.Module):
|
||||
"""ConvTranspose1d with some builtin handling of asymmetric or causal padding
|
||||
and normalization.
|
||||
"""
|
||||
def __init__(self, in_channels: int, out_channels: int,
|
||||
kernel_size: int, stride: int = 1, causal: bool = False,
|
||||
norm: str = 'none', trim_right_ratio: float = 1.,
|
||||
norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
|
||||
super().__init__()
|
||||
self.convtr = NormConvTranspose1d(in_channels, out_channels, kernel_size, stride,
|
||||
causal=causal, norm=norm, norm_kwargs=norm_kwargs)
|
||||
self.causal = causal
|
||||
self.trim_right_ratio = trim_right_ratio
|
||||
assert self.causal or self.trim_right_ratio == 1., \
|
||||
"`trim_right_ratio` != 1.0 only makes sense for causal convolutions"
|
||||
assert self.trim_right_ratio >= 0. and self.trim_right_ratio <= 1.
|
||||
|
||||
def forward(self, x):
|
||||
kernel_size = self.convtr.convtr.kernel_size[0]
|
||||
stride = self.convtr.convtr.stride[0]
|
||||
padding_total = kernel_size - stride
|
||||
|
||||
y = self.convtr(x)
|
||||
|
||||
# We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be
|
||||
# removed at the very end, when keeping only the right length for the output,
|
||||
# as removing it here would require also passing the length at the matching layer
|
||||
# in the encoder.
|
||||
if self.causal:
|
||||
# Trim the padding on the right according to the specified ratio
|
||||
# if trim_right_ratio = 1.0, trim everything from right
|
||||
padding_right = math.ceil(padding_total * self.trim_right_ratio)
|
||||
padding_left = padding_total - padding_right
|
||||
y = unpad1d(y, (padding_left, padding_right))
|
||||
else:
|
||||
# Asymmetric padding required for odd strides
|
||||
padding_right = padding_total // 2
|
||||
padding_left = padding_total - padding_right
|
||||
y = unpad1d(y, (padding_left, padding_right))
|
||||
return y
|
||||
|
||||
class SLSTM(nn.Module):
|
||||
"""
|
||||
LSTM without worrying about the hidden state, nor the layout of the data.
|
||||
Expects input as convolutional layout.
|
||||
"""
|
||||
def __init__(self, dimension: int, num_layers: int = 2, skip: bool = True):
|
||||
super().__init__()
|
||||
self.skip = skip
|
||||
self.lstm = nn.LSTM(dimension, dimension, num_layers)
|
||||
self.hidden = None
|
||||
|
||||
def forward(self, x):
|
||||
x = x.permute(2, 0, 1)
|
||||
if self.training:
|
||||
y, _ = self.lstm(x)
|
||||
else:
|
||||
y, self.hidden = self.lstm(x, self.hidden)
|
||||
if self.skip:
|
||||
y = y + x
|
||||
y = y.permute(1, 2, 0)
|
||||
return y
|
||||
Reference in New Issue
Block a user