Update model.py

model.py

@@ -1,49 +1,42 @@
 import torch
 import torch.nn as nn
-from .positional_encoding import FramePositionalEncoding
-from .config import TTSConfig
+from .positional_encoding import PositionalEncoding
+
 
 class RealtimeTTS(nn.Module):
-    def __init__(self, config=TTSConfig()):
+    def __init__(self, config):
         super().__init__()
 
-        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
-
-        self.encoder = nn.TransformerEncoder(
-            nn.TransformerEncoderLayer(
-                d_model=config.d_model,
-                nhead=config.n_heads,
-                batch_first=True
-            ),
-            num_layers=config.num_encoder_layers
+        self.embedding = nn.Embedding(
+            config.vocab_size,
+            config.d_model
         )
 
-        self.frame_pe = FramePositionalEncoding(config.d_model)
-
-        self.decoder = nn.TransformerDecoder(
-            nn.TransformerDecoderLayer(
-                d_model=config.d_model,
-                nhead=config.n_heads,
-                batch_first=True
-            ),
-            num_layers=config.num_decoder_layers
+        self.positional_encoding = PositionalEncoding(
+            config.d_model,
+            config.max_seq_len
         )
 
-        self.mel_projection = nn.Linear(config.d_model, config.mel_bins)
-
-    def forward(self, text_tokens, mel_inputs):
-        # Text embedding
-        x = self.embedding(text_tokens)
-
-        # Text encoding
-        memory = self.encoder(x)
-
-        # Frame positional encoding
-        mel_inputs = self.frame_pe(mel_inputs)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=config.d_model,
+            nhead=config.nhead,
+            dim_feedforward=config.dim_feedforward,
+            batch_first=True
+        )
 
-        # Decode mel frames
-        out = self.decoder(mel_inputs, memory)
+        self.transformer = nn.TransformerEncoder(
+            encoder_layer,
+            num_layers=config.num_layers
+        )
 
-        mel_output = self.mel_projection(out)
+        self.output_linear = nn.Linear(
+            config.d_model,
+            80  # mel bins
+        )
 
-        return mel_output
+    def forward(self, tokens, mel_input):
+        x = self.embedding(tokens)
+        x = self.positional_encoding(x)
+        x = self.transformer(x)
+        mel = self.output_linear(x)
+        return mel