initial commit

README.md

@@ -1,7 +1,7 @@
 ---
 title: Translation Stack
 emoji: 🏢
-colorFrom: purple
+colorFrom: blue
 colorTo: green
 sdk: gradio
 sdk_version: 5.33.1

app.py

@@ -0,0 +1,573 @@
+import os
+import re
+import time
+import random
+import numpy as np
+import math
+import shutil
+# import base64 # Not directly needed for Gradio filepath output
+
+# 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.nn.functional as F
+import torchaudio
+import librosa
+# import librosa.display # Not used in pipeline
+
+# Text and Audio Processing
+from unidecode import unidecode
+# from inflect import engine # Not explicitly used in pipeline, consider removing
+# import pydub # Not explicitly used in pipeline, consider removing
+import soundfile as sf
+
+# Transformers
+from transformers import (
+    WhisperProcessor, WhisperForConditionalGeneration,
+    MarianTokenizer, MarianMTModel,
+)
+from huggingface_hub import hf_hub_download
+
+# Gradio and Hugging Face Spaces
+import gradio as gr
+import spaces # <<< --- ADD THIS IMPORT --- <<<
+
+# --- Global Configuration & Device Setup ---
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"--- Initializing on device: {DEVICE} ---") # This will run when the Space builds/starts
+
+# --- Part 1: TTS Model Components (Your Custom TTS) ---
+# ... (Keep all your Hyperparams, text_to_seq, audio processing for TTS, and Model class definitions:
+# EncoderBlock, DecoderBlock, EncoderPreNet, PostNet, DecoderPreNet, TransformerTTS)
+# ... (Ensure TransformerTTS and its sub-modules are correctly defined as in your previous code)
+# --- (Start of your model definitions - make sure this is complete from your previous code) ---
+class Hyperparams:
+  seed = 42
+  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',
+    't', 'u', 'v', 'w', 'x', 'y', 'z', 'à', 'â', 'è', 'é', 'ê', 'ü',
+    '’', '“', '”'
+  ]
+  sr = 22050
+  n_fft = 2048
+  n_stft = int((n_fft//2) + 1)
+  hop_length = int(n_fft/8.0)
+  win_length = int(n_fft/2.0)
+  mel_freq = 128
+  max_mel_time = 1024
+  power = 2.0
+  text_num_embeddings = 2*len(symbols)
+  embedding_size = 256
+  encoder_embedding_size = 512
+  dim_feedforward = 1024
+  postnet_embedding_size = 1024
+  encoder_kernel_size = 3
+  postnet_kernel_size = 5
+  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()
+
+symbol_to_id = {s: i for i, s in enumerate(hp.symbols)}
+def text_to_seq(text):
+  text = text.lower()
+  text = unidecode(text)
+  seq = []
+  for s in text:
+    _id = symbol_to_id.get(s, None)
+    if _id is not None:
+      seq.append(_id)
+  seq.append(symbol_to_id["EOS"])
+  return torch.IntTensor(seq)
+
+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)
+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).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 = 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:
+    ones = sequence_lengths.new_ones(sequence_lengths.size(0), max_length)
+    range_tensor = ones.cumsum(dim=1)
+    return sequence_lengths.unsqueeze(1) >= range_tensor
+
+class EncoderBlock(nn.Module): # Your EncoderBlock definition
+    def __init__(self):
+        super(EncoderBlock, self).__init__()
+        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(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(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): # Your DecoderBlock definition
+    def __init__(self):
+        super(DecoderBlock, self).__init__()
+        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(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(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(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): # Your EncoderPreNet definition
+    def __init__(self):
+        super(EncoderPreNet, self).__init__()
+        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, 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): # Your PostNet definition
+    def __init__(self):
+        super(PostNet, self).__init__()
+        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 # Store original for residual connection if postnet predicts residual
+        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) # No Tanh on last layer for mel usually
+        x = x.transpose(1, 2)
+        return x # This is the residual, should be added to original mel_linear
+
+class DecoderPreNet(nn.Module): # Your DecoderPreNet definition
+    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=self.training)
+        x = self.linear_2(x)
+        x = F.relu(x)
+        x = F.dropout(x, p=0.5, training=self.training)
+        return x
+
+class TransformerTTS(nn.Module): # Your TransformerTTS definition
+    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(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) # Stop token
+        self.norm_memory = nn.LayerNorm(normalized_shape=hp.embedding_size)
+        self.device = device
+
+    def forward(self, text, text_len, mel, mel_len): # For training/teacher-forcing
+        # ... (Your detailed forward pass for training, with all masks)
+        N = text.shape[0]; S = text.shape[1]; TIME = mel.shape[1]
+        current_device = text.device
+
+        src_key_padding_mask = torch.zeros((N, S), device=current_device, dtype=torch.bool).masked_fill(~mask_from_seq_lengths(text_len, max_length=S), True)
+        src_mask = None # Typically encoder self-attention doesn't use a causal mask
+
+        tgt_key_padding_mask = torch.zeros((N, TIME), device=current_device, dtype=torch.bool).masked_fill(~mask_from_seq_lengths(mel_len, max_length=TIME), True)
+        tgt_mask = torch.zeros((TIME, TIME), device=current_device).masked_fill(torch.triu(torch.full((TIME, TIME), True, device=current_device, dtype=torch.bool), diagonal=1), float("-inf"))
+        memory_mask = None # Cross-attention mask, typically not needed unless specific structure
+
+        text_x = self.encoder_prenet(text)
+        pos_codes = self.pos_encoding(torch.arange(hp.max_mel_time, device=current_device))
+        text_s_dim = text_x.shape[1]
+        text_x = text_x + pos_codes[:text_s_dim]
+
+        text_x = self.encoder_block_1(text_x, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)
+        text_x = self.encoder_block_2(text_x, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)
+        text_x = self.encoder_block_3(text_x, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)
+        memory = self.norm_memory(text_x)
+
+        mel_x = self.decoder_prenet(mel)
+        mel_time_dim = mel_x.shape[1]
+        mel_x = mel_x + pos_codes[:mel_time_dim]
+
+        mel_x = self.decoder_block_1(x=mel_x, memory=memory, x_attn_mask=tgt_mask, x_key_padding_mask=tgt_key_padding_mask, memory_attn_mask=memory_mask, memory_key_padding_mask=src_key_padding_mask)
+        mel_x = self.decoder_block_2(x=mel_x, memory=memory, x_attn_mask=tgt_mask, x_key_padding_mask=tgt_key_padding_mask, memory_attn_mask=memory_mask, memory_key_padding_mask=src_key_padding_mask)
+        mel_x = self.decoder_block_3(x=mel_x, memory=memory, x_attn_mask=tgt_mask, x_key_padding_mask=tgt_key_padding_mask, memory_attn_mask=memory_mask, memory_key_padding_mask=src_key_padding_mask)
+
+        mel_linear = self.linear_1(mel_x)
+        mel_postnet_residual = self.postnet(mel_linear) # Postnet predicts residual
+        mel_postnet = mel_linear + mel_postnet_residual
+
+        stop_token = self.linear_2(mel_x) # Sigmoid applied later
+
+        # Masking for training outputs
+        bool_mel_mask = tgt_key_padding_mask.unsqueeze(-1).repeat(1, 1, hp.mel_freq)
+        mel_linear = mel_linear.masked_fill(bool_mel_mask, 0.0)
+        mel_postnet = mel_postnet.masked_fill(bool_mel_mask, 0.0)
+        # Ensure stop_token is [N, TIME]
+        stop_token = stop_token.masked_fill(tgt_key_padding_mask.unsqueeze(-1) if stop_token.dim() == 3 else tgt_key_padding_mask, 1e3)
+        if stop_token.dim() == 3 and stop_token.shape[2] == 1:
+            stop_token = stop_token.squeeze(-1)
+
+
+        return mel_postnet, mel_linear, stop_token
+
+
+    @torch.no_grad()
+    def inference(self, text, max_length=800, stop_token_threshold=0.5): # text: [1, seq_len]
+        self.eval()
+        N = text.shape[0] # Should be 1
+        current_device = text.device
+        text_lengths = torch.tensor([text.shape[1]], device=current_device)
+
+        # Encoder pass (once)
+        src_key_padding_mask_inf = torch.zeros((N, text.shape[1]), device=current_device, dtype=torch.bool) # All False initially
+        # No, src_key_padding_mask should be based on actual text length, even if N=1, S=text.shape[1]
+        # For inference with single item, it's often all False (no padding in input text usually)
+        # However, to be consistent with how `mask_from_seq_lengths` works:
+        src_key_padding_mask_inf = ~mask_from_seq_lengths(text_lengths, text.shape[1])
+
+
+        encoder_output = self.encoder_prenet(text)
+        pos_codes = self.pos_encoding(torch.arange(hp.max_mel_time, device=current_device))
+        text_s_dim = encoder_output.shape[1]
+        encoder_output = encoder_output + pos_codes[:text_s_dim]
+
+        encoder_output = self.encoder_block_1(encoder_output, key_padding_mask=src_key_padding_mask_inf)
+        encoder_output = self.encoder_block_2(encoder_output, key_padding_mask=src_key_padding_mask_inf)
+        encoder_output = self.encoder_block_3(encoder_output, key_padding_mask=src_key_padding_mask_inf)
+        memory = self.norm_memory(encoder_output)
+
+        # Decoder pass (iterative)
+        mel_input = torch.zeros((N, 1, hp.mel_freq), device=current_device) # SOS frame
+        generated_mel_frames = []
+
+        for i in range(max_length):
+            mel_lengths_inf = torch.tensor([mel_input.shape[1]], device=current_device)
+            # For decoder self-attention, causal mask is needed
+            tgt_mask_inf = torch.zeros((mel_input.shape[1], mel_input.shape[1]), device=current_device).masked_fill(
+                torch.triu(torch.full((mel_input.shape[1], mel_input.shape[1]), True, device=current_device, dtype=torch.bool), diagonal=1), float("-inf")
+            )
+            # Decoder input padding mask (all False as we build it frame by frame, no padding yet)
+            tgt_key_padding_mask_inf = torch.zeros((N, mel_input.shape[1]), device=current_device, dtype=torch.bool)
+
+
+            mel_x = self.decoder_prenet(mel_input)
+            mel_time_dim = mel_input.shape[1]
+            mel_x = mel_x + pos_codes[:mel_time_dim] # Positional encoding for current mel sequence
+
+            mel_x = self.decoder_block_1(x=mel_x, memory=memory, x_attn_mask=tgt_mask_inf, x_key_padding_mask=tgt_key_padding_mask_inf, memory_key_padding_mask=src_key_padding_mask_inf)
+            mel_x = self.decoder_block_2(x=mel_x, memory=memory, x_attn_mask=tgt_mask_inf, x_key_padding_mask=tgt_key_padding_mask_inf, memory_key_padding_mask=src_key_padding_mask_inf)
+            mel_x = self.decoder_block_3(x=mel_x, memory=memory, x_attn_mask=tgt_mask_inf, x_key_padding_mask=tgt_key_padding_mask_inf, memory_key_padding_mask=src_key_padding_mask_inf)
+
+            mel_linear_step = self.linear_1(mel_x[:, -1:, :]) # Predict only for the last frame
+            mel_postnet_residual_step = self.postnet(mel_linear_step)
+            current_mel_frame = mel_linear_step + mel_postnet_residual_step
+
+            generated_mel_frames.append(current_mel_frame)
+            mel_input = torch.cat([mel_input, current_mel_frame], dim=1) # Append to input for next step
+
+            # Stop token prediction (based on the last decoder output before linear to mel)
+            stop_token_logit = self.linear_2(mel_x[:, -1:, :]) # Stop token from last frame's decoder hidden state
+            stop_token_prob = torch.sigmoid(stop_token_logit.squeeze())
+
+            if stop_token_prob > stop_token_threshold:
+                # print(f"Stop token threshold reached at step {i+1}")
+                break
+            if mel_input.shape[1] > hp.max_mel_time -1: # Safety break based on max_mel_time
+                # print(f"Max mel time {hp.max_mel_time} almost reached.")
+                break
+
+
+        if not generated_mel_frames:
+            print("Warning: TTS inference produced no mel frames.")
+            return torch.zeros((N, 0, hp.mel_freq), device=current_device) # Return empty tensor
+
+        final_mel_output = torch.cat(generated_mel_frames, dim=1)
+        return final_mel_output # Removed stop_token_outputs as it's not used by caller
+# --- (End of your model definitions) ---
+
+# --- Part 2: Model Loading ---
+# (Same as before - ensure TTS_MODEL = TransformerTTS(device=DEVICE).to(DEVICE) is used)
+TTS_MODEL_HUB_ID = "MoHamdyy/transformer-tts-ljspeech"
+ASR_HUB_ID       = "MoHamdyy/whisper-stt-model"
+MARIAN_HUB_ID    = "MoHamdyy/marian-ar-en-translation"
+
+TTS_MODEL = None
+stt_processor = None
+stt_model = None
+mt_tokenizer = None
+mt_model = None
+
+# Wrap model loading in a function to clearly see when it happens or to potentially delay it.
+# For Spaces, global loading is fine and preferred as it happens once.
+print("--- Starting Model Loading ---")
+try:
+    print(f"Loading TTS model ({TTS_MODEL_HUB_ID}) to {DEVICE}...")
+    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) # Load to target device directly
+    TTS_MODEL = TransformerTTS(device=DEVICE).to(DEVICE)
+    model_state_dict = state.get("model", state.get("state_dict", state))
+    TTS_MODEL.load_state_dict(model_state_dict)
+    TTS_MODEL.eval()
+    print("TTS model loaded successfully.")
+except Exception as e:
+    print(f"Error loading TTS model: {e}")
+
+try:
+    print(f"Loading STT (Whisper) model ({ASR_HUB_ID}) to {DEVICE}...")
+    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}")
+
+try:
+    print(f"Loading TTT (MarianMT) model ({MARIAN_HUB_ID}) to {DEVICE}...")
+    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}")
+print("--- Model Loading Complete ---")
+
+
+# --- Part 3: Full Pipeline Function for Gradio ---
+@spaces.GPU # <<< --- APPLY THE DECORATOR HERE --- <<<
+def full_speech_translation_pipeline_gradio(audio_input_path):
+    # This print will show the device context *inside* the decorated function
+    # For ZeroGPU, this should ideally show 'cuda:X' when the function is executed
+    current_processing_device = next(stt_model.parameters()).device if stt_model else "CPU (STT model not loaded)"
+    print(f"--- @spaces.GPU function: Pipeline running on device: {current_processing_device} ---")
+
+
+    if not all([TTS_MODEL, stt_processor, stt_model, mt_tokenizer, mt_model]):
+        error_msg = "Critical Error: One or more models failed to load during Space initialization. Cannot process."
+        print(error_msg)
+        # Raising gr.Error is better for UI feedback
+        raise gr.Error(error_msg)
+
+
+    if audio_input_path is None:
+        # This case should ideally be handled by Gradio's input validation or a check before calling.
+        # If it still occurs, provide a clear message.
+        raise gr.Error("No audio file provided. Please upload an audio file.")
+
+    print(f"--- GRADIO PIPELINE START (GPU context): Processing {audio_input_path} ---")
+
+    # STT Stage
+    arabic_transcript = "STT Error: Processing failed."
+    try:
+        print("STT: Loading and resampling audio...")
+        wav, sr = torchaudio.load(audio_input_path)
+        if wav.size(0) > 1: wav = wav.mean(dim=0, keepdim=True)
+        target_sr_stt = stt_processor.feature_extractor.sampling_rate
+        if sr != target_sr_stt: wav = torchaudio.transforms.Resample(sr, target_sr_stt)(wav)
+        # Move wav to the STT model's device *before* converting to numpy if STT model is on GPU
+        audio_array_stt = wav.to(DEVICE).squeeze().cpu().numpy() # Process on DEVICE, then to CPU for numpy
+
+        print("STT: Extracting features and transcribing...")
+        # Ensure inputs are on the same device as the model
+        inputs_stt = stt_processor(audio_array_stt, sampling_rate=target_sr_stt, return_tensors="pt").input_features.to(DEVICE)
+        forced_ids = stt_processor.get_decoder_prompt_ids(language="arabic", task="transcribe")
+        with torch.no_grad():
+            generated_ids = stt_model.generate(inputs_stt, forced_decoder_ids=forced_ids, max_new_tokens=256)
+        arabic_transcript = stt_processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+        print(f"STT Output: {arabic_transcript}")
+    except Exception as e:
+        print(f"STT Error: {e}")
+        raise gr.Error(f"STT processing failed: {e}")
+
+
+    # TTT Stage
+    english_translation = "TTT Error: Processing failed."
+    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, truncation=True).to(DEVICE)
+            with torch.no_grad():
+                translated_ids = mt_model.generate(**batch, max_length=512) # max_new_tokens can also be used
+            english_translation = mt_tokenizer.batch_decode(translated_ids, skip_special_tokens=True)[0].strip()
+            print(f"TTT Output: {english_translation}")
+        except Exception as e:
+            print(f"TTT Error: {e}")
+            raise gr.Error(f"TTT processing failed: {e}")
+
+    else:
+        if not arabic_transcript or arabic_transcript.startswith("STT Error"):
+            english_translation = "(Skipped TTT due to STT failure or empty STT output)"
+        print(english_translation)
+
+
+    # TTS Stage
+    output_tts_audio_filepath = None
+    if english_translation and not english_translation.startswith("TTT Error") and TTS_MODEL:
+        try:
+            print("TTS: Synthesizing English speech...")
+            if not english_translation.strip():
+                 print("TTS Warning: Empty string for synthesis. Skipping TTS.")
+            else:
+                sequence = text_to_seq(english_translation).unsqueeze(0).to(DEVICE)
+                # max_length for TTS inference refers to max output mel frames
+                generated_mel = TTS_MODEL.inference(sequence, max_length=hp.max_mel_time - 50, stop_token_threshold=0.5)
+
+                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 :
+                    # TTS model's inverse_mel_spec_to_wav expects mel on DEVICE and returns wav on CPU
+                    # The mel from inference should be [N, mel_len, mel_bins]
+                    # inverse_mel_spec_to_wav might expect [mel_bins, mel_len]
+                    mel_for_vocoder = generated_mel.detach().squeeze(0).transpose(0, 1) # to [mel_len, mel_bins]
+                    audio_tensor = inverse_mel_spec_to_wav(mel_for_vocoder) # This function handles .to(DEVICE) internally
+                    synthesized_audio_np = audio_tensor.cpu().numpy() # Ensure output is on CPU for soundfile
+                    print(f"TTS: Synthesized audio shape: {synthesized_audio_np.shape}")
+
+                    timestamp = int(time.time()*1000) # more unique
+                    output_tts_audio_filepath = f"output_audio_{timestamp}.wav"
+                    sf.write(output_tts_audio_filepath, synthesized_audio_np, hp.sr)
+                    print(f"TTS: Synthesized audio saved to: {output_tts_audio_filepath}")
+                else:
+                    print("TTS Warning: Generated mel spectrogram was empty or invalid.")
+        except Exception as e:
+            print(f"TTS Error: {e}")
+            # Do not raise gr.Error here if a partial result (text) is still useful
+            # output_tts_audio_filepath will remain None
+            english_translation += f" (TTS Error: {e})" # Append error to text
+    else:
+        if not TTS_MODEL: print("TTS SKIPPED: Model not loaded.")
+        elif not (english_translation and not english_translation.startswith("TTT Error")):
+             print("TTS SKIPPED: (Due to TTT failure or empty TTT output)")
+
+
+    print(f"--- GRADIO PIPELINE END (GPU context) ---")
+    return arabic_transcript, english_translation, output_tts_audio_filepath
+
+
+# --- Part 4: Gradio Interface Definition ---
+# (Same as before)
+iface = gr.Interface(
+    fn=full_speech_translation_pipeline_gradio,
+    inputs=[
+        gr.Audio(type="filepath", label="Upload Arabic Speech")
+    ],
+    outputs=[
+        gr.Textbox(label="Arabic Transcript (STT)"),
+        gr.Textbox(label="English Translation (TTT)"),
+        gr.Audio(label="Synthesized English Speech (TTS)", type="filepath")
+    ],
+    title="Arabic to English Speech Translation (ZeroGPU)",
+    description="Upload an Arabic audio file. Transcribed to Arabic (Whisper), translated to English (MarianMT), synthesized to English speech (Custom TransformerTTS).",
+    allow_flagging="never",
+    # examples=[["sample.wav"]] # If you add a sample.wav to your repo
+)
+
+# --- Part 5: Launch for Spaces (and local testing) ---
+if __name__ == '__main__':
+    # Clean up temp audio files from previous local runs
+    for f_name in os.listdir("."):
+        if f_name.startswith("output_audio_") and f_name.endswith(".wav"):
+            try:
+                os.remove(f_name)
+            except OSError:
+                pass # Ignore if file is already gone or locked
+    print("Starting Gradio interface locally with debug mode...")
+    iface.launch(debug=True, share=False) # share=False for local, Spaces handles public URL

requirements.txt

@@ -0,0 +1,24 @@
+# Machine Learning & Audio
+torch
+torchaudio
+# For CPU, ensure correct PyTorch version if not using ZeroGPU:
+# torch --index-url https://download.pytorch.org/whl/cpu
+# For CUDA (ZeroGPU is CUDA-based, HF Spaces will handle this if ZeroGPU is selected)
+# torch --index-url https://download.pytorch.org/whl/cu118 
+# (Check latest recommended CUDA version for ZeroGPU on HF docs)
+
+transformers
+librosa
+soundfile
+pydub
+unidecode
+inflect
+huggingface_hub
+sentencepiece # Often needed by tokenizers
+
+# Gradio
+gradio >=4.0.0 # Use a recent version of Gradio
+
+# Other utilities
+numpy
+pandas # Though pandas is not explicitly used in the pipeline, it's in your imports