|
from torch import nn |
|
|
|
from TTS.tts.layers.generic.pos_encoding import PositionalEncoding |
|
from TTS.tts.layers.generic.transformer import FFTransformerBlock |
|
|
|
|
|
class DurationPredictor(nn.Module): |
|
def __init__(self, num_chars, hidden_channels, hidden_channels_ffn, num_heads): |
|
super().__init__() |
|
self.embed = nn.Embedding(num_chars, hidden_channels) |
|
self.pos_enc = PositionalEncoding(hidden_channels, dropout_p=0.1) |
|
self.FFT = FFTransformerBlock(hidden_channels, num_heads, hidden_channels_ffn, 2, 0.1) |
|
self.out_layer = nn.Conv1d(hidden_channels, 1, 1) |
|
|
|
def forward(self, text, text_lengths): |
|
|
|
emb = self.embed(text) |
|
emb = self.pos_enc(emb.transpose(1, 2)) |
|
x = self.FFT(emb, text_lengths) |
|
x = self.out_layer(x).squeeze(-1) |
|
return x |
|
|
