app.py

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