Gabriel de Marmiesse
c90fc8ced6
The french 6 layers variant wasn't included because it's not good enough yet. Use the `french_24l` variant for now.
52 lines
1.4 KiB
Python
52 lines
1.4 KiB
Python
import torch
|
|
from torch import nn
|
|
|
|
from pocket_tts.modules.conv import StreamingConv1d, StreamingConvTranspose1d
|
|
|
|
|
|
class ConvDownsample1d(nn.Module):
|
|
"""
|
|
Downsampling by some integer amount `stride` using convolutions
|
|
with a kernel size of twice the stride.
|
|
"""
|
|
|
|
def __init__(self, stride: int, dimension: int, out_dimension: int | None = None):
|
|
super().__init__()
|
|
if out_dimension is None:
|
|
out_dimension = dimension
|
|
|
|
self.conv = StreamingConv1d(
|
|
dimension,
|
|
out_dimension,
|
|
kernel_size=2 * stride,
|
|
stride=stride,
|
|
groups=1,
|
|
bias=False,
|
|
pad_mode="replicate",
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor, model_state: dict | None):
|
|
return self.conv(x, model_state)
|
|
|
|
|
|
class ConvTrUpsample1d(nn.Module):
|
|
"""
|
|
Upsample by some integer amount `stride` using transposed convolutions.
|
|
"""
|
|
|
|
def __init__(self, stride: int, dimension: int, in_dimension: int | None = None):
|
|
super().__init__()
|
|
if in_dimension is None:
|
|
in_dimension = dimension
|
|
self.convtr = StreamingConvTranspose1d(
|
|
in_dimension,
|
|
dimension,
|
|
kernel_size=2 * stride,
|
|
stride=stride,
|
|
groups=dimension,
|
|
bias=False,
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor, model_state: dict | None):
|
|
return self.convtr(x, model_state)
|