changed app to gradio

app.py

@@ -3,26 +3,26 @@ import re
 import time
 import random
 import numpy as np
-import pandas as pd # Keep if hp or other parts use it, though not directly in pipeline
-import math # Keep if hp or other parts use it
-# import shutil # Not needed for Gradio file handling
-# import base64 # Not needed for Gradio audio output
+import pandas as pd
+import math
+import shutil
+import base64
 
 # Torch and Audio
 import torch
 import torch.nn as nn
-# import torch.optim as optim # Not needed for inference
-# from torch.utils.data import Dataset, DataLoader # Not needed for inference
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
 import torch.nn.functional as F
 import torchaudio
-# import librosa # Not strictly needed if not plotting in Gradio
-# import librosa.display # Not strictly needed if not plotting in Gradio
+import librosa
+import librosa.display
 
 # Text and Audio Processing
 from unidecode import unidecode
-from inflect import engine as inflect_engine_tts # Renamed to avoid conflict
-# import pydub # Not needed for Gradio audio output
-# import soundfile as sf # Gradio handles audio output directly
+from inflect import engine
+import pydub
+import soundfile as sf
 
 # Transformers
 from transformers import (
@@ -30,27 +30,22 @@ from transformers import (
     MarianTokenizer, MarianMTModel,
 )
 
-# Gradio
-import gradio as gr
-from huggingface_hub import hf_hub_download # For downloading models
+# API Server
+from fastapi import FastAPI, UploadFile, File
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles # <--- ADD THIS IMPORT
+
 
-# --- Configuration & Device ---
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using device: {DEVICE}")
 
-torch.manual_seed(42)
-np.random.seed(42)
-random.seed(42)
-if torch.cuda.is_available():
-    torch.cuda.manual_seed_all(42)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
+# Part 2: TTS Model Components (from your notebook)
 
-# --- Hyperparams Class (VERBATIM from your notebook) ---
+
+# Hyperparameters
 class Hyperparams:
   seed = 42
-  csv_path = "path/to/metadata.csv" # Not used directly
-  wav_path = "path/to/wavs" # Not used directly
+  # We won't use these dataset paths, but keep them for hp object integrity
+  csv_path = "path/to/metadata.csv"
+  wav_path = "path/to/wavs"
   symbols = [
     'EOS', ' ', '!', ',', '-', '.', ';', '?', 'a', 'b', 'c', 'd', 'e', 'f', 
     'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 
@@ -64,7 +59,7 @@ class Hyperparams:
   win_length = int(n_fft/2.0)
   mel_freq = 128
   max_mel_time = 1024
-  power = 2.0 # For spec_transform if used, not directly by inverse_mel
+  power = 2.0
   text_num_embeddings = 2*len(symbols)  
   embedding_size = 256
   encoder_embedding_size = 512 
@@ -72,16 +67,17 @@ class Hyperparams:
   postnet_embedding_size = 1024
   encoder_kernel_size = 3
   postnet_kernel_size = 5
-  ampl_multiplier = 10.0 # For pow_to_db_mel_spec
-  ampl_amin = 1e-10 # For pow_to_db_mel_spec
-  db_multiplier = 1.0 # For pow_to_db_mel_spec
-  ampl_ref = 1.0 # For db_to_power_mel_spec
-  ampl_power = 1.0 # For db_to_power_mel_spec
-  max_db = 100 # For pow_to_db_mel_spec
-  scale_db = 10 # For pow_to_db_mel_spec & db_to_power_mel_spec
+  ampl_multiplier = 10.0
+  ampl_amin = 1e-10
+  db_multiplier = 1.0
+  ampl_ref = 1.0
+  ampl_power = 1.0
+  max_db = 100
+  scale_db = 10
+
 hp = Hyperparams()
 
-# --- TTS Text & Audio Processing (VERBATIM from your notebook) ---
+# Text to Sequence
 symbol_to_id = {s: i for i, s in enumerate(hp.symbols)}
 def text_to_seq(text):
   text = text.lower()
@@ -93,19 +89,27 @@ def text_to_seq(text):
   seq.append(symbol_to_id["EOS"])
   return torch.IntTensor(seq)
 
+# Audio Processing
+spec_transform = torchaudio.transforms.Spectrogram(n_fft=hp.n_fft, win_length=hp.win_length, hop_length=hp.hop_length, power=hp.power)
+mel_scale_transform = torchaudio.transforms.MelScale(n_mels=hp.mel_freq, sample_rate=hp.sr, n_stft=hp.n_stft)
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 mel_inverse_transform = torchaudio.transforms.InverseMelScale(n_mels=hp.mel_freq, sample_rate=hp.sr, n_stft=hp.n_stft).to(DEVICE)
-griffnlim_transform = torchaudio.transforms.GriffinLim(n_fft=hp.n_fft, win_length=hp.win_length, hop_length=hp.hop_length, power=1.0).to(DEVICE) # Explicit power=1.0 for magnitude
+griffnlim_transform = torchaudio.transforms.GriffinLim(n_fft=hp.n_fft, win_length=hp.win_length, hop_length=hp.hop_length).to(DEVICE)
+
+def pow_to_db_mel_spec(mel_spec):
+  mel_spec = torchaudio.functional.amplitude_to_DB(mel_spec, multiplier=hp.ampl_multiplier, amin=hp.ampl_amin, db_multiplier=hp.db_multiplier, top_db=hp.max_db)
+  mel_spec = mel_spec/hp.scale_db
+  return mel_spec
 
 def db_to_power_mel_spec(mel_spec):
-  mel_spec_scaled = mel_spec * hp.scale_db # Corrected: use a different variable name
-  mel_spec_amp = torchaudio.functional.DB_to_amplitude(mel_spec_scaled, ref=hp.ampl_ref, power=hp.ampl_power)  
-  return mel_spec_amp
-
-def inverse_mel_spec_to_wav(mel_spec): # Expects [Freq, Time]
-  mel_spec_on_device = mel_spec.to(DEVICE)
-  power_mel_spec = db_to_power_mel_spec(mel_spec_on_device) # This is amplitude
-  spectrogram = mel_inverse_transform(power_mel_spec) # Amplitude mel to linear amplitude
-  pseudo_wav = griffnlim_transform(spectrogram) # Linear amplitude to wav
+  mel_spec = mel_spec*hp.scale_db
+  mel_spec = torchaudio.functional.DB_to_amplitude(mel_spec, ref=hp.ampl_ref, power=hp.ampl_power)  
+  return mel_spec
+
+def inverse_mel_spec_to_wav(mel_spec):
+  power_mel_spec = db_to_power_mel_spec(mel_spec.to(DEVICE))
+  spectrogram = mel_inverse_transform(power_mel_spec)
+  pseudo_wav = griffnlim_transform(spectrogram)
   return pseudo_wav
 
 def mask_from_seq_lengths(sequence_lengths: torch.Tensor, max_length: int) -> torch.BoolTensor:
@@ -113,266 +117,287 @@ def mask_from_seq_lengths(sequence_lengths: torch.Tensor, max_length: int) -> to
     range_tensor = ones.cumsum(dim=1)
     return sequence_lengths.unsqueeze(1) >= range_tensor
     
-# --- TransformerTTS Model Architecture (VERBATIM from your FastAPI code) ---
-class EncoderBlock(nn.Module): # VERBATIM
+# --- TransformerTTS Model Architecture (Copied from notebook)
+class EncoderBlock(nn.Module):
     def __init__(self):
         super(EncoderBlock, self).__init__()
-        self.norm_1 = nn.LayerNorm(hp.embedding_size)
-        self.attn = torch.nn.MultiheadAttention(hp.embedding_size, 4, dropout=0.1, batch_first=True)
+        self.norm_1 = nn.LayerNorm(normalized_shape=hp.embedding_size)
+        self.attn = torch.nn.MultiheadAttention(embed_dim=hp.embedding_size, num_heads=4, dropout=0.1, batch_first=True)
         self.dropout_1 = torch.nn.Dropout(0.1)
-        self.norm_2 = nn.LayerNorm(hp.embedding_size)
+        self.norm_2 = nn.LayerNorm(normalized_shape=hp.embedding_size)
         self.linear_1 = nn.Linear(hp.embedding_size, hp.dim_feedforward)
         self.dropout_2 = torch.nn.Dropout(0.1)
         self.linear_2 = nn.Linear(hp.dim_feedforward, hp.embedding_size)
         self.dropout_3 = torch.nn.Dropout(0.1)
     def forward(self, x, attn_mask=None, key_padding_mask=None):
-        x_out = self.norm_1(x); x_out, _ = self.attn(x_out, x_out, x_out, attn_mask=attn_mask, key_padding_mask=key_padding_mask)
-        x_out = self.dropout_1(x_out); x = x + x_out    
-        x_out = self.norm_2(x) ; x_out = self.linear_1(x_out); x_out = F.relu(x_out); x_out = self.dropout_2(x_out)
-        x_out = self.linear_2(x_out); x_out = self.dropout_3(x_out)
-        x = x + x_out; return x
-
-class DecoderBlock(nn.Module): # VERBATIM
+        x_out = self.norm_1(x)
+        x_out, _ = self.attn(query=x_out, key=x_out, value=x_out, attn_mask=attn_mask, key_padding_mask=key_padding_mask)
+        x_out = self.dropout_1(x_out)
+        x = x + x_out    
+        x_out = self.norm_2(x) 
+        x_out = self.linear_1(x_out)
+        x_out = F.relu(x_out)
+        x_out = self.dropout_2(x_out)
+        x_out = self.linear_2(x_out)
+        x_out = self.dropout_3(x_out)
+        x = x + x_out
+        return x
+
+class DecoderBlock(nn.Module):
     def __init__(self):
         super(DecoderBlock, self).__init__()
-        self.norm_1 = nn.LayerNorm(hp.embedding_size)
-        self.self_attn = torch.nn.MultiheadAttention(hp.embedding_size, 4, dropout=0.1, batch_first=True)
+        self.norm_1 = nn.LayerNorm(normalized_shape=hp.embedding_size)
+        self.self_attn = torch.nn.MultiheadAttention(embed_dim=hp.embedding_size, num_heads=4, dropout=0.1, batch_first=True)
         self.dropout_1 = torch.nn.Dropout(0.1)
-        self.norm_2 = nn.LayerNorm(hp.embedding_size)
-        self.attn = torch.nn.MultiheadAttention(hp.embedding_size, 4, dropout=0.1, batch_first=True)    
+        self.norm_2 = nn.LayerNorm(normalized_shape=hp.embedding_size)
+        self.attn = torch.nn.MultiheadAttention(embed_dim=hp.embedding_size, num_heads=4, dropout=0.1, batch_first=True)    
         self.dropout_2 = torch.nn.Dropout(0.1)
-        self.norm_3 = nn.LayerNorm(hp.embedding_size)
+        self.norm_3 = nn.LayerNorm(normalized_shape=hp.embedding_size)
         self.linear_1 = nn.Linear(hp.embedding_size, hp.dim_feedforward)
         self.dropout_3 = torch.nn.Dropout(0.1)
         self.linear_2 = nn.Linear(hp.dim_feedforward, hp.embedding_size)
         self.dropout_4 = torch.nn.Dropout(0.1)
     def forward(self, x, memory, x_attn_mask=None, x_key_padding_mask=None, memory_attn_mask=None, memory_key_padding_mask=None):
-        x_out, _ = self.self_attn(x, x, x, attn_mask=x_attn_mask, key_padding_mask=x_key_padding_mask)
-        x_out = self.dropout_1(x_out); x = self.norm_1(x + x_out)
-        x_out, _ = self.attn(x, memory, memory, attn_mask=memory_attn_mask, key_padding_mask=memory_key_padding_mask)
-        x_out = self.dropout_2(x_out); x = self.norm_2(x + x_out)
-        x_out = self.linear_1(x); x_out = F.relu(x_out); x_out = self.dropout_3(x_out)
-        x_out = self.linear_2(x_out); x_out = self.dropout_4(x_out)
-        x = self.norm_3(x + x_out); return x
-
-class EncoderPreNet(nn.Module): # VERBATIM
+        x_out, _ = self.self_attn(query=x, key=x, value=x, attn_mask=x_attn_mask, key_padding_mask=x_key_padding_mask)
+        x_out = self.dropout_1(x_out)
+        x = self.norm_1(x + x_out)
+        x_out, _ = self.attn(query=x, key=memory, value=memory, attn_mask=memory_attn_mask, key_padding_mask=memory_key_padding_mask)
+        x_out = self.dropout_2(x_out)
+        x = self.norm_2(x + x_out)
+        x_out = self.linear_1(x)
+        x_out = F.relu(x_out)
+        x_out = self.dropout_3(x_out)
+        x_out = self.linear_2(x_out)
+        x_out = self.dropout_4(x_out)
+        x = self.norm_3(x + x_out)
+        return x
+
+class EncoderPreNet(nn.Module):
     def __init__(self):
         super(EncoderPreNet, self).__init__()
-        self.embedding = nn.Embedding(hp.text_num_embeddings, hp.encoder_embedding_size)
+        self.embedding = nn.Embedding(num_embeddings=hp.text_num_embeddings, embedding_dim=hp.encoder_embedding_size)
         self.linear_1 = nn.Linear(hp.encoder_embedding_size, hp.encoder_embedding_size)
         self.linear_2 = nn.Linear(hp.encoder_embedding_size, hp.embedding_size)
-        self.conv_1 = nn.Conv1d(hp.encoder_embedding_size, hp.encoder_embedding_size, hp.encoder_kernel_size, 1, int((hp.encoder_kernel_size-1)/2),1)
-        self.bn_1 = nn.BatchNorm1d(hp.encoder_embedding_size); self.dropout_1 = nn.Dropout(0.5)
-        self.conv_2 = nn.Conv1d(hp.encoder_embedding_size, hp.encoder_embedding_size, hp.encoder_kernel_size, 1, int((hp.encoder_kernel_size-1)/2),1)
-        self.bn_2 = nn.BatchNorm1d(hp.encoder_embedding_size); self.dropout_2 = nn.Dropout(0.5)
-        self.conv_3 = nn.Conv1d(hp.encoder_embedding_size, hp.encoder_embedding_size, hp.encoder_kernel_size, 1, int((hp.encoder_kernel_size-1)/2),1)
-        self.bn_3 = nn.BatchNorm1d(hp.encoder_embedding_size); self.dropout_3 = nn.Dropout(0.5)    
+        self.conv_1 = nn.Conv1d(hp.encoder_embedding_size, hp.encoder_embedding_size, kernel_size=hp.encoder_kernel_size, stride=1, padding=int((hp.encoder_kernel_size - 1) / 2), dilation=1)
+        self.bn_1 = nn.BatchNorm1d(hp.encoder_embedding_size)
+        self.dropout_1 = torch.nn.Dropout(0.5)
+        self.conv_2 = nn.Conv1d(hp.encoder_embedding_size, hp.encoder_embedding_size, kernel_size=hp.encoder_kernel_size, stride=1, padding=int((hp.encoder_kernel_size - 1) / 2), dilation=1)
+        self.bn_2 = nn.BatchNorm1d(hp.encoder_embedding_size)
+        self.dropout_2 = torch.nn.Dropout(0.5)
+        self.conv_3 = nn.Conv1d(hp.encoder_embedding_size, hp.encoder_embedding_size, kernel_size=hp.encoder_kernel_size, stride=1, padding=int((hp.encoder_kernel_size - 1) / 2), dilation=1)
+        self.bn_3 = nn.BatchNorm1d(hp.encoder_embedding_size)
+        self.dropout_3 = torch.nn.Dropout(0.5)    
     def forward(self, text):
-        x = self.embedding(text); x = self.linear_1(x); x = x.transpose(2,1) 
-        x = self.conv_1(x); x = self.bn_1(x); x = F.relu(x); x = self.dropout_1(x)
-        x = self.conv_2(x); x = self.bn_2(x); x = F.relu(x); x = self.dropout_2(x)
-        x = self.conv_3(x); x = self.bn_3(x); x = F.relu(x); x = self.dropout_3(x)
-        x = x.transpose(1,2); x = self.linear_2(x); return x
-
-class PostNet(nn.Module): # VERBATIM
+        x = self.embedding(text)
+        x = self.linear_1(x)
+        x = x.transpose(2, 1)
+        x = self.conv_1(x)
+        x = self.bn_1(x)
+        x = F.relu(x)
+        x = self.dropout_1(x)
+        x = self.conv_2(x)
+        x = self.bn_2(x)
+        x = F.relu(x)
+        x = self.dropout_2(x)
+        x = self.conv_3(x)
+        x = self.bn_3(x)    
+        x = F.relu(x)
+        x = self.dropout_3(x)
+        x = x.transpose(1, 2)
+        x = self.linear_2(x)
+        return x
+
+class PostNet(nn.Module):
     def __init__(self):
         super(PostNet, self).__init__()
-        self.conv_1 = nn.Conv1d(hp.mel_freq, hp.postnet_embedding_size, hp.postnet_kernel_size, 1, int((hp.postnet_kernel_size-1)/2),1)
-        self.bn_1 = nn.BatchNorm1d(hp.postnet_embedding_size); self.dropout_1 = nn.Dropout(0.5)
-        self.conv_2 = nn.Conv1d(hp.postnet_embedding_size, hp.postnet_embedding_size, hp.postnet_kernel_size, 1, int((hp.postnet_kernel_size-1)/2),1)
-        self.bn_2 = nn.BatchNorm1d(hp.postnet_embedding_size); self.dropout_2 = nn.Dropout(0.5)
-        self.conv_3 = nn.Conv1d(hp.postnet_embedding_size, hp.postnet_embedding_size, hp.postnet_kernel_size, 1, int((hp.postnet_kernel_size-1)/2),1)
-        self.bn_3 = nn.BatchNorm1d(hp.postnet_embedding_size); self.dropout_3 = nn.Dropout(0.5)
-        self.conv_4 = nn.Conv1d(hp.postnet_embedding_size, hp.postnet_embedding_size, hp.postnet_kernel_size, 1, int((hp.postnet_kernel_size-1)/2),1)
-        self.bn_4 = nn.BatchNorm1d(hp.postnet_embedding_size); self.dropout_4 = nn.Dropout(0.5)
-        self.conv_5 = nn.Conv1d(hp.postnet_embedding_size, hp.postnet_embedding_size, hp.postnet_kernel_size, 1, int((hp.postnet_kernel_size-1)/2),1)
-        self.bn_5 = nn.BatchNorm1d(hp.postnet_embedding_size); self.dropout_5 = nn.Dropout(0.5)
-        self.conv_6 = nn.Conv1d(hp.postnet_embedding_size, hp.mel_freq, hp.postnet_kernel_size, 1, int((hp.postnet_kernel_size-1)/2),1)
-        self.bn_6 = nn.BatchNorm1d(hp.mel_freq); self.dropout_6 = nn.Dropout(0.5)
+        self.conv_1 = nn.Conv1d(hp.mel_freq, hp.postnet_embedding_size, kernel_size=hp.postnet_kernel_size, stride=1, padding=int((hp.postnet_kernel_size - 1) / 2), dilation=1)
+        self.bn_1 = nn.BatchNorm1d(hp.postnet_embedding_size)
+        self.dropout_1 = torch.nn.Dropout(0.5)
+        self.conv_2 = nn.Conv1d(hp.postnet_embedding_size, hp.postnet_embedding_size, kernel_size=hp.postnet_kernel_size, stride=1, padding=int((hp.postnet_kernel_size - 1) / 2), dilation=1)
+        self.bn_2 = nn.BatchNorm1d(hp.postnet_embedding_size)
+        self.dropout_2 = torch.nn.Dropout(0.5)
+        self.conv_3 = nn.Conv1d(hp.postnet_embedding_size, hp.postnet_embedding_size, kernel_size=hp.postnet_kernel_size, stride=1, padding=int((hp.postnet_kernel_size - 1) / 2), dilation=1)
+        self.bn_3 = nn.BatchNorm1d(hp.postnet_embedding_size)
+        self.dropout_3 = torch.nn.Dropout(0.5)
+        self.conv_4 = nn.Conv1d(hp.postnet_embedding_size, hp.postnet_embedding_size, kernel_size=hp.postnet_kernel_size, stride=1, padding=int((hp.postnet_kernel_size - 1) / 2), dilation=1)
+        self.bn_4 = nn.BatchNorm1d(hp.postnet_embedding_size)
+        self.dropout_4 = torch.nn.Dropout(0.5)
+        self.conv_5 = nn.Conv1d(hp.postnet_embedding_size, hp.postnet_embedding_size, kernel_size=hp.postnet_kernel_size, stride=1, padding=int((hp.postnet_kernel_size - 1) / 2), dilation=1)
+        self.bn_5 = nn.BatchNorm1d(hp.postnet_embedding_size)
+        self.dropout_5 = torch.nn.Dropout(0.5)
+        self.conv_6 = nn.Conv1d(hp.postnet_embedding_size, hp.mel_freq, kernel_size=hp.postnet_kernel_size, stride=1, padding=int((hp.postnet_kernel_size - 1) / 2), dilation=1)
+        self.bn_6 = nn.BatchNorm1d(hp.mel_freq)
+        self.dropout_6 = torch.nn.Dropout(0.5)
     def forward(self, x):
-        x_orig = x; x = x.transpose(2,1)
-        x = self.conv_1(x); x = self.bn_1(x); x = torch.tanh(x); x = self.dropout_1(x)
-        x = self.conv_2(x); x = self.bn_2(x); x = torch.tanh(x); x = self.dropout_2(x)
-        x = self.conv_3(x); x = self.bn_3(x); x = torch.tanh(x); x = self.dropout_3(x)
-        x = self.conv_4(x); x = self.bn_4(x); x = torch.tanh(x); x = self.dropout_4(x)
-        x = self.conv_5(x); x = self.bn_5(x); x = torch.tanh(x); x = self.dropout_5(x)
-        x = self.conv_6(x); x = self.bn_6(x); x = self.dropout_6(x)
-        x = x.transpose(1,2); return x # Original postnet in repo is residual, added in TransformerTTS.forward
-
-class DecoderPreNet(nn.Module): # VERBATIM
+        x = x.transpose(2, 1)
+        x = self.conv_1(x)
+        x = self.bn_1(x); x = torch.tanh(x); x = self.dropout_1(x)
+        x = self.conv_2(x)
+        x = self.bn_2(x); x = torch.tanh(x); x = self.dropout_2(x)
+        x = self.conv_3(x)
+        x = self.bn_3(x); x = torch.tanh(x); x = self.dropout_3(x)
+        x = self.conv_4(x)
+        x = self.bn_4(x); x = torch.tanh(x); x = self.dropout_4(x)
+        x = self.conv_5(x)
+        x = self.bn_5(x); x = torch.tanh(x); x = self.dropout_5(x)
+        x = self.conv_6(x)
+        x = self.bn_6(x); x = self.dropout_6(x)
+        x = x.transpose(1, 2)
+        return x
+
+class DecoderPreNet(nn.Module):
     def __init__(self):
         super(DecoderPreNet, self).__init__()
         self.linear_1 = nn.Linear(hp.mel_freq, hp.embedding_size)
         self.linear_2 = nn.Linear(hp.embedding_size, hp.embedding_size)
     def forward(self, x):
-        x = self.linear_1(x); x = F.relu(x)
-        x = F.dropout(x, p=0.5, training=True) # Dropout always on
-        x = self.linear_2(x); x = F.relu(x)    
-        x = F.dropout(x, p=0.5, training=True) # Dropout always on
+        x = self.linear_1(x)
+        x = F.relu(x)
+        x = F.dropout(x, p=0.5, training=True)
+        x = self.linear_2(x)
+        x = F.relu(x)    
+        x = F.dropout(x, p=0.5, training=True)
         return x    
 
-class TransformerTTS(nn.Module): # VERBATIM (init had device=DEVICE, now model is moved after init)
-    def __init__(self): # Removed device=DEVICE from here
+class TransformerTTS(nn.Module):
+    def __init__(self, device=DEVICE):
         super(TransformerTTS, self).__init__()
         self.encoder_prenet = EncoderPreNet()
         self.decoder_prenet = DecoderPreNet()
         self.postnet = PostNet()
-        self.pos_encoding = nn.Embedding(hp.max_mel_time, hp.embedding_size)
-        self.encoder_block_1 = EncoderBlock(); self.encoder_block_2 = EncoderBlock(); self.encoder_block_3 = EncoderBlock()
-        self.decoder_block_1 = DecoderBlock(); self.decoder_block_2 = DecoderBlock(); self.decoder_block_3 = DecoderBlock()
+        self.pos_encoding = nn.Embedding(num_embeddings=hp.max_mel_time, embedding_dim=hp.embedding_size)
+        self.encoder_block_1 = EncoderBlock()
+        self.encoder_block_2 = EncoderBlock()
+        self.encoder_block_3 = EncoderBlock()
+        self.decoder_block_1 = DecoderBlock()
+        self.decoder_block_2 = DecoderBlock()
+        self.decoder_block_3 = DecoderBlock()
         self.linear_1 = nn.Linear(hp.embedding_size, hp.mel_freq) 
         self.linear_2 = nn.Linear(hp.embedding_size, 1)
-        self.norm_memory = nn.LayerNorm(hp.embedding_size)
-        # Mask attributes will be set in forward pass, as per your code
-        self.src_key_padding_mask = None; self.src_mask = None
-        self.tgt_key_padding_mask = None; self.tgt_mask = None; self.memory_mask = None
-
-    def forward(self, text, text_len, mel, mel_len): # VERBATIM
-        N = text.shape[0]; S_text_in = text.shape[1]; TIME_mel_in = mel.shape[1]
-        current_device = text.device
-
-        self.src_key_padding_mask = torch.zeros((N, S_text_in), device=current_device).masked_fill(~mask_from_seq_lengths(text_len, max_length=S_text_in), float("-inf"))
-        self.src_mask = torch.zeros((S_text_in, S_text_in), device=current_device).masked_fill(torch.triu(torch.full((S_text_in, S_text_in), True, dtype=torch.bool, device=current_device), diagonal=1), float("-inf"))
-        self.tgt_key_padding_mask = torch.zeros((N, TIME_mel_in), device=current_device).masked_fill(~mask_from_seq_lengths(mel_len, max_length=TIME_mel_in), float("-inf"))
-        self.tgt_mask = torch.zeros((TIME_mel_in, TIME_mel_in), device=current_device).masked_fill(torch.triu(torch.full((TIME_mel_in, TIME_mel_in), True, device=current_device, dtype=torch.bool), diagonal=1), float("-inf"))
-        self.memory_mask = torch.zeros((TIME_mel_in, S_text_in), device=current_device).masked_fill(torch.triu(torch.full((TIME_mel_in, S_text_in), True, device=current_device, dtype=torch.bool), diagonal=1), float("-inf"))    
-        
+        self.norm_memory = nn.LayerNorm(normalized_shape=hp.embedding_size)
+    def forward(self, text, text_len, mel, mel_len):  
+        N = text.shape[0]; S = text.shape[1]; TIME = mel.shape[1]
+        self.src_key_padding_mask = torch.zeros((N, S), device=text.device).masked_fill(~mask_from_seq_lengths(text_len, max_length=S), float("-inf"))
+        self.src_mask = torch.zeros((S, S), device=text.device).masked_fill(torch.triu(torch.full((S, S), True, dtype=torch.bool), diagonal=1).to(text.device), float("-inf"))
+        self.tgt_key_padding_mask = torch.zeros((N, TIME), device=mel.device).masked_fill(~mask_from_seq_lengths(mel_len, max_length=TIME), float("-inf"))
+        self.tgt_mask = torch.zeros((TIME, TIME), device=mel.device).masked_fill(torch.triu(torch.full((TIME, TIME), True, device=mel.device, dtype=torch.bool), diagonal=1), float("-inf"))
+        self.memory_mask = torch.zeros((TIME, S), device=mel.device).masked_fill(torch.triu(torch.full((TIME, S), True, device=mel.device, dtype=torch.bool), diagonal=1), float("-inf"))    
         text_x = self.encoder_prenet(text) 
-        pos_codes = self.pos_encoding(torch.arange(hp.max_mel_time, device=current_device))
-        S_text_processed = text_x.shape[1]; text_x = text_x + pos_codes[:S_text_processed] # Use actual S after prenet
-        
+        pos_codes = self.pos_encoding(torch.arange(hp.max_mel_time).to(mel.device))
+        S = text_x.shape[1]; text_x = text_x + pos_codes[:S]
         text_x = self.encoder_block_1(text_x, attn_mask = self.src_mask, key_padding_mask = self.src_key_padding_mask)
         text_x = self.encoder_block_2(text_x, attn_mask = self.src_mask, key_padding_mask = self.src_key_padding_mask)    
         text_x = self.encoder_block_3(text_x, attn_mask = self.src_mask, key_padding_mask = self.src_key_padding_mask)
         text_x = self.norm_memory(text_x)
-        
-        mel_x = self.decoder_prenet(mel); 
-        TIME_mel_processed = mel_x.shape[1]; mel_x = mel_x + pos_codes[:TIME_mel_processed] # Use actual T after prenet
-
+        mel_x = self.decoder_prenet(mel); mel_x = mel_x + pos_codes[:TIME]
         mel_x = self.decoder_block_1(x=mel_x, memory=text_x, x_attn_mask=self.tgt_mask, x_key_padding_mask=self.tgt_key_padding_mask, memory_attn_mask=self.memory_mask, memory_key_padding_mask=self.src_key_padding_mask)
         mel_x = self.decoder_block_2(x=mel_x, memory=text_x, x_attn_mask=self.tgt_mask, x_key_padding_mask=self.tgt_key_padding_mask, memory_attn_mask=self.memory_mask, memory_key_padding_mask=self.src_key_padding_mask)
         mel_x = self.decoder_block_3(x=mel_x, memory=text_x, x_attn_mask=self.tgt_mask, x_key_padding_mask=self.tgt_key_padding_mask, memory_attn_mask=self.memory_mask, memory_key_padding_mask=self.src_key_padding_mask)
-        
         mel_linear = self.linear_1(mel_x)
-        postnet_residual_out = self.postnet(mel_linear) # PostNet output
-        mel_postnet = mel_linear + postnet_residual_out # Add residual
-        stop_token = self.linear_2(mel_x) # (N, TIME, 1)
-
-        # Masking output based on padding
-        # self.tgt_key_padding_mask is -inf for padded, 0 for unpadded.
-        # .ne(0) makes it True for padded, False for unpadded. This is correct for masked_fill.
-        bool_mel_padding_mask = self.tgt_key_padding_mask.ne(0) 
-        
-        mel_linear = mel_linear.masked_fill(bool_mel_padding_mask.unsqueeze(-1).expand_as(mel_linear), 0)
-        mel_postnet = mel_postnet.masked_fill(bool_mel_padding_mask.unsqueeze(-1).expand_as(mel_postnet), 0)
-        stop_token = stop_token.masked_fill(bool_mel_padding_mask.unsqueeze(-1).expand_as(stop_token), 1e3).squeeze(2)
+        mel_postnet = self.postnet(mel_linear)
+        mel_postnet = mel_linear + mel_postnet
+        stop_token = self.linear_2(mel_x)
+        bool_mel_mask = self.tgt_key_padding_mask.ne(0).unsqueeze(-1).repeat(1, 1, hp.mel_freq)
+        mel_linear = mel_linear.masked_fill(bool_mel_mask, 0)
+        mel_postnet = mel_postnet.masked_fill(bool_mel_mask, 0)
+        stop_token = stop_token.masked_fill(bool_mel_mask[:, :, 0].unsqueeze(-1), 1e3).squeeze(2)
         return mel_postnet, mel_linear, stop_token 
 
-    @torch.no_grad() # VERBATIM from your FastAPI code (with .item() fix)
-    def inference(self, text, max_length=800, stop_token_threshold=0.5, with_tqdm=False): # with_tqdm was False in pipeline call
-        self.eval(); self.train(False) # As per your original
-        model_device = next(self.parameters()).device
-        
-        text_on_device = text.to(model_device)
-        text_lengths = torch.tensor([text_on_device.shape[1]],dtype=torch.long).unsqueeze(0).to(model_device) # Ensure text_lengths is also 2D [1,1] or [1]
-        
+    @torch.no_grad()
+    def inference(self, text, max_length=800, stop_token_threshold=0.5, with_tqdm=True):
+        self.eval(); self.train(False)
+        text_lengths = torch.tensor(text.shape[1]).unsqueeze(0).to(DEVICE)
         N = 1
-        SOS = torch.zeros((N, 1, hp.mel_freq), device=model_device)
+        SOS = torch.zeros((N, 1, hp.mel_freq), device=DEVICE)
         mel_padded = SOS
-        mel_lengths = torch.tensor([1],dtype=torch.long).unsqueeze(0).to(model_device) # Ensure mel_lengths is also 2D [1,1] or [1]
-        
-        stop_token_outputs = torch.FloatTensor([]).to(model_device) # text.device might be CPU if text wasn't on device
-        
-        # Use local tqdm to avoid conflict if tqdm is imported elsewhere
-        from tqdm import tqdm as tqdm_local 
-        iters = tqdm_local(range(max_length), desc="TTS Inference") if with_tqdm else range(max_length)
-        
-        final_mel_postnet_output = SOS # To store the output from the last forward pass
-
+        mel_lengths = torch.tensor(1).unsqueeze(0).to(DEVICE)
+        stop_token_outputs = torch.FloatTensor([]).to(text.device)
+        iters = range(max_length)
         for _ in iters:
-            # mel_postnet is (N, T_current_input_len, Freq)
-            mel_postnet, mel_linear, stop_token = self(text_on_device, text_lengths, mel_padded, mel_lengths)
-            final_mel_postnet_output = mel_postnet # This is the full sequence predicted in this step
-
-            # Append last frame of mel_postnet for next input step
+            mel_postnet, mel_linear, stop_token = self(text, text_lengths, mel_padded, mel_lengths)
             mel_padded = torch.cat([mel_padded, mel_postnet[:, -1:, :]], dim=1)
-            mel_lengths = torch.tensor([mel_padded.shape[1]],dtype=torch.long).unsqueeze(0).to(model_device)
-            
-            # stop_token is (N, T_current_input_len)
-            # Check stop condition for the last frame of the input sequence
-            if (torch.sigmoid(stop_token[:, -1].squeeze()) > stop_token_threshold).item():      
+            if torch.sigmoid(stop_token[:, -1]) > stop_token_threshold:      
                 break
             else:
-                # stop_token[:, -1:] is (N, 1)
                 stop_token_outputs = torch.cat([stop_token_outputs, stop_token[:, -1:]], dim=1)
+                mel_lengths = torch.tensor(mel_padded.shape[1]).unsqueeze(0).to(DEVICE)
+        return mel_postnet, stop_token_outputs
         
-        # final_mel_postnet_output contains SOS. Strip it.
-        if final_mel_postnet_output.shape[1] > 1: # If more than just SOS frame
-            mel_to_return = final_mel_postnet_output[:, 1:, :] 
-        else: # Only SOS was processed, or nothing
-            mel_to_return = torch.empty((N, 0, hp.mel_freq), device=model_device)
-            if mel_to_return.shape[1] == 0: # ensure stop_token_outputs is also empty
-                 stop_token_outputs = torch.empty_like(stop_token_outputs[:,:0])
+# Part 3: Model Loading
 
 
-        return mel_to_return, stop_token_outputs
-        
-# --- Part 3: Model Loading (from Hugging Face Hub - VERBATIM from your FastAPI code) ---
+# IMPORTANT: These paths assume you have placed the downloaded models
+# into a 'models' subfolder in your project directory.
+# ---------------------------------
+# --- Part 3: Model Loading (from Hugging Face Hub)
+# ---------------------------------
+
+# IMPORTANT: Replace "your-username" with your Hugging Face username
+# and the model names with the ones you created on the Hub.
 TTS_MODEL_HUB_ID = "MoHamdyy/transformer-tts-ljspeech"
 ASR_HUB_ID       = "MoHamdyy/whisper-stt-model"
 MARIAN_HUB_ID    = "MoHamdyy/marian-ar-en-translation"
 
+# Helper function to download the TTS model file from the Hub
+from huggingface_hub import hf_hub_download
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 print("Loading models from Hugging Face Hub to device:", DEVICE)
-TTS_MODEL = None; stt_processor = None; stt_model = None; mt_tokenizer = None; mt_model = None
 
+# Load TTS Model from Hub
 try:
     print("Loading TTS model...")
+    # Download the .pt file from its repo
     tts_model_path = hf_hub_download(repo_id=TTS_MODEL_HUB_ID, filename="train_SimpleTransfromerTTS.pt")
     state = torch.load(tts_model_path, map_location=DEVICE)
-    TTS_MODEL = TransformerTTS().to(DEVICE) # Create instance then move to DEVICE
-    if "model" in state: TTS_MODEL.load_state_dict(state["model"])
-    elif "state_dict" in state: TTS_MODEL.load_state_dict(state["state_dict"])
-    else: TTS_MODEL.load_state_dict(state)
+    TTS_MODEL = TransformerTTS().to(DEVICE)
+    # Check for the correct key in the state dictionary
+    if "model" in state:
+        TTS_MODEL.load_state_dict(state["model"])
+    elif "state_dict" in state:
+        TTS_MODEL.load_state_dict(state["state_dict"])
+    else:
+        TTS_MODEL.load_state_dict(state) # Assume the whole file is the state_dict
     TTS_MODEL.eval()
     print("TTS model loaded successfully.")
-except Exception as e: print(f"Error loading TTS model: {e}")
+except Exception as e:
+    print(f"Error loading TTS model: {e}")
+    TTS_MODEL = None
 
+# Load STT (Whisper) Model from Hub
 try:
     print("Loading STT (Whisper) model...")
     stt_processor = WhisperProcessor.from_pretrained(ASR_HUB_ID)
     stt_model = WhisperForConditionalGeneration.from_pretrained(ASR_HUB_ID).to(DEVICE).eval()
     print("STT model loaded successfully.")
-except Exception as e: print(f"Error loading STT model: {e}")
+except Exception as e:
+    print(f"Error loading STT model: {e}")
+    stt_processor = None
+    stt_model = None
 
+# Load TTT (MarianMT) Model from Hub
 try:
     print("Loading TTT (MarianMT) model...")
     mt_tokenizer = MarianTokenizer.from_pretrained(MARIAN_HUB_ID)
     mt_model = MarianMTModel.from_pretrained(MARIAN_HUB_ID).to(DEVICE).eval()
     print("TTT model loaded successfully.")
-except Exception as e: print(f"Error loading TTT model: {e}")
+except Exception as e:
+    print(f"Error loading TTT model: {e}")
+    mt_tokenizer = None
+    mt_model = None
+
+
 
-# --- Part 4: Full Pipeline Function (VERBATIM from your FastAPI code, adapted for Gradio output) ---
-def full_speech_translation_pipeline_gradio(audio_input_path: str): # Renamed for clarity
+# Part 4: Full Pipeline Function
+
+
+def full_speech_translation_pipeline(audio_input_path: str):
     print(f"--- PIPELINE START: Processing {audio_input_path} ---")
-    # Check if models are loaded
-    if not all([stt_processor, stt_model, mt_tokenizer, mt_model, TTS_MODEL]):
-        error_msg = "One or more models failed to load. Please check logs."
-        print(error_msg)
-        return error_msg, error_msg, (hp.sr, np.array([]).astype(np.float32))
-
-    if audio_input_path is None: # Gradio provides a path for uploaded/recorded audio
-        msg = "Error: No audio input received by Gradio."
-        print(msg)
-        return msg, "", (hp.sr, np.array([]).astype(np.float32))
-    
-    if not os.path.exists(audio_input_path):
-        # This case might happen if Gradio passes a temp path that gets cleaned up too quickly,
-        # or if there's an issue with how Gradio handles file paths.
-        # For Gradio `type="filepath"`, the path should be valid.
-        msg = f"Error: Audio file path provided by Gradio does not exist: {audio_input_path}"
+    if audio_input_path is None or not os.path.exists(audio_input_path):
+        msg = "Error: Audio file not provided or not found."
         print(msg)
-        return msg, "", (hp.sr, np.array([]).astype(np.float32))
-
+        # Return empty/default values
+        return "Error: No file", "", (hp.sr, np.array([]).astype(np.float32))
 
     # STT Stage
     arabic_transcript = "STT Error: Processing failed."
@@ -391,16 +416,12 @@ def full_speech_translation_pipeline_gradio(audio_input_path: str): # Renamed fo
             generated_ids = stt_model.generate(inputs, forced_decoder_ids=forced_ids, max_length=448)
         arabic_transcript = stt_processor.decode(generated_ids[0], skip_special_tokens=True).strip()
         print(f"STT Output: {arabic_transcript}")
-        if not arabic_transcript: arabic_transcript = "(STT: No speech detected or empty transcript)"
     except Exception as e:
-        print(f"STT Error: {e}"); import traceback; traceback.print_exc()
-        arabic_transcript = f"STT Error: {e}"
-
+        print(f"STT Error: {e}")
 
     # TTT Stage
     english_translation = "TTT Error: Processing failed."
-    tts_status_message = "" # For appending TTS status to English text
-    if arabic_transcript and not arabic_transcript.startswith("STT Error") and not arabic_transcript.startswith("(STT:"):
+    if arabic_transcript and not arabic_transcript.startswith("STT Error"):
         try:
             print("TTT: Translating to English...")
             batch = mt_tokenizer(arabic_transcript, return_tensors="pt", padding=True).to(DEVICE)
@@ -408,74 +429,72 @@ def full_speech_translation_pipeline_gradio(audio_input_path: str): # Renamed fo
                 translated_ids = mt_model.generate(**batch, max_length=512)
             english_translation = mt_tokenizer.batch_decode(translated_ids, skip_special_tokens=True)[0].strip()
             print(f"TTT Output: {english_translation}")
-            if not english_translation: english_translation = "(TTT: Empty translation)"
         except Exception as e:
-            print(f"TTT Error: {e}"); import traceback; traceback.print_exc()
-            english_translation = f"TTT Error: {e}"
-    elif arabic_transcript.startswith("STT Error") or arabic_transcript.startswith("(STT:"):
-        english_translation = "(Skipped TTT due to STT issue)"
+            print(f"TTT Error: {e}")
+    else:
+        english_translation = "(Skipped TTT due to STT failure)"
         print(english_translation)
-    else: # Should not happen if STT produces some output
-        english_translation = "(Skipped TTT: Unknown STT state)"
-
 
     # TTS Stage
     synthesized_audio_np = np.array([]).astype(np.float32)
-    if english_translation and not english_translation.startswith("TTT Error") and not english_translation.startswith("(Skipped") and not english_translation.startswith("(TTT:"):
+    if english_translation and not english_translation.startswith("TTT Error"):
         try:
             print("TTS: Synthesizing English speech...")
-            sequence = text_to_seq(english_translation).unsqueeze(0).to(DEVICE) # Ensure input is on TTS_MODEL's device
-            
-            # Make sure TTS_MODEL is on the correct device before inference
-            TTS_MODEL.to(DEVICE) # Redundant if already done, but safe
-            TTS_MODEL.eval()     # Ensure eval mode
-
+            sequence = text_to_seq(english_translation).unsqueeze(0).to(DEVICE)
             generated_mel, _ = TTS_MODEL.inference(sequence, max_length=hp.max_mel_time-20, stop_token_threshold=0.5, with_tqdm=False)
             
             print(f"TTS: Generated mel shape: {generated_mel.shape if generated_mel is not None else 'None'}")
-            if generated_mel is not None and generated_mel.numel() > 0 and generated_mel.shape[1] > 0: # Check if time dimension has frames
-                mel_for_vocoder = generated_mel.detach().squeeze(0).transpose(0, 1) # [F, T]
+            if generated_mel is not None and generated_mel.numel() > 0:
+                mel_for_vocoder = generated_mel.detach().squeeze(0).transpose(0, 1)
                 audio_tensor = inverse_mel_spec_to_wav(mel_for_vocoder)
                 synthesized_audio_np = audio_tensor.cpu().numpy()
                 print(f"TTS: Synthesized audio shape: {synthesized_audio_np.shape}")
-            else:
-                tts_status_message = "(TTS Error: Empty mel generated)"
-                print(tts_status_message)
         except Exception as e:
-            print(f"TTS Error: {e}"); import traceback; traceback.print_exc()
-            tts_status_message = f"(TTS Error: {e})"
-    elif english_translation.startswith("TTT Error") or english_translation.startswith("(Skipped") or english_translation.startswith("(TTT:"):
-        tts_status_message = "(Skipped TTS due to TTT/Input issue)"
-    else: # Should not happen if TTT produces some output
-        tts_status_message = "(Skipped TTS: Unknown TTT state)"
-        
+            print(f"TTS Error: {e}")
+    
     print(f"--- PIPELINE END ---")
-    # Combine English translation with any TTS status message
-    final_english_display = english_translation
-    if tts_status_message:
-        final_english_display += f" {tts_status_message}"
-        
-    return arabic_transcript, final_english_display.strip(), (hp.sr, synthesized_audio_np)
-
-# --- Part 5: Gradio Interface ---
-print("Setting up Gradio interface...")
-demo = gr.Interface(
-    fn=full_speech_translation_pipeline_gradio,
-    inputs=gr.Audio(sources=["microphone", "upload"], type="filepath", label="Arabic Speech Input"),
-    outputs=[
-        gr.Textbox(label="Arabic Transcript (STT)"),
-        gr.Textbox(label="English Translation & TTS Status"),
-        gr.Audio(label="Synthesized English Speech (TTS)", type="numpy") # type="numpy" expects (sr, data)
-    ],
-    title="Arabic Speech-to-Text -> Translation -> English Text-to-Speech",
-    description="Upload an Arabic audio file or record from your microphone. The system will transcribe it to Arabic, translate it to English, and then synthesize the English text into audible speech.",
-    allow_flagging="never",
-    examples=[["/kaggle/input/testtt/test_audio.ogg"]] if os.path.exists("/kaggle/input/testtt/test_audio.ogg") else None # Optional example
+    return arabic_transcript, english_translation, (hp.sr, synthesized_audio_np)
+
+
+# Part 5: FastAPI Application
+
+app = FastAPI()
+
+# Allow Cross-Origin Resource Sharing (CORS) for your frontend
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Allows all origins
+    allow_credentials=True,
+    allow_methods=["*"],  # Allows all methods
+    allow_headers=["*"],  # Allows all headers
 )
-demo.launch(debug=True)
-# if __name__ == "__main__":
-#     print("Launching Gradio app...")
-#     # When running on Hugging Face Spaces, HF handles the launch.
-#     # For local testing, you might need a specific host/port.
-#     # HF Spaces will look for a `demo.launch()` or `iface.launch()`
-#     # demo.launch(debug=True) # debug=True for more detailed Gradio logs
+
+@app.post("/process-speech/")
+async def create_upload_file(file: UploadFile = File(...)):
+    # Save the uploaded file temporarily
+    temp_path = f"/tmp/{file.filename}"
+    with open(temp_path, "wb") as buffer:
+        shutil.copyfileobj(file.file, buffer)
+    
+    # Run the full pipeline
+    arabic, english, (sr, audio_np) = full_speech_translation_pipeline(temp_path)
+    
+    # Prepare the audio to be sent back as base64
+    audio_base64 = ""
+    if audio_np.size > 0:
+        temp_wav_path = "/tmp/output.wav"
+        sf.write(temp_wav_path, audio_np, sr)
+        with open(temp_wav_path, "rb") as wav_file:
+            audio_bytes = wav_file.read()
+            audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
+    
+    # Return all results in a single JSON response
+    return {
+        "arabic_transcript": arabic,
+        "english_translation": english,
+        "audio_data": {
+            "sample_rate": sr,
+            "base64": audio_base64
+        }
+    }
+app.mount("/", StaticFiles(directory="static", html=True), name="static")