Spaces:
Sleeping
Sleeping
| import os | |
| 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 | |
| # 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 # Not strictly needed if not plotting in Gradio | |
| # import librosa.display # Not strictly needed if not plotting in Gradio | |
| # 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 | |
| # Transformers | |
| from transformers import ( | |
| WhisperProcessor, WhisperForConditionalGeneration, | |
| MarianTokenizer, MarianMTModel, | |
| ) | |
| # Gradio | |
| import gradio as gr | |
| from huggingface_hub import hf_hub_download # For downloading models | |
| # --- 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 | |
| # --- Hyperparams Class (VERBATIM from your notebook) --- | |
| class Hyperparams: | |
| seed = 42 | |
| csv_path = "path/to/metadata.csv" # Not used directly | |
| wav_path = "path/to/wavs" # Not used directly | |
| 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 # For spec_transform if used, not directly by inverse_mel | |
| 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 # 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 | |
| hp = Hyperparams() | |
| # --- TTS Text & Audio Processing (VERBATIM from your notebook) --- | |
| symbol_to_id = {s: i for i, s in enumerate(hp.symbols)} | |
| def text_to_seq(text): | |
| text = text.lower() | |
| 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) | |
| 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 | |
| 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 | |
| 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 | |
| # --- TransformerTTS Model Architecture (VERBATIM from your FastAPI code) --- | |
| class EncoderBlock(nn.Module): # VERBATIM | |
| 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.dropout_1 = torch.nn.Dropout(0.1) | |
| self.norm_2 = nn.LayerNorm(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 | |
| 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.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.dropout_2 = torch.nn.Dropout(0.1) | |
| self.norm_3 = nn.LayerNorm(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 | |
| def __init__(self): | |
| super(EncoderPreNet, self).__init__() | |
| self.embedding = nn.Embedding(hp.text_num_embeddings, 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) | |
| 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 | |
| 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) | |
| 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 | |
| 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 | |
| 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 | |
| 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.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")) | |
| 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 | |
| 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_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) | |
| return mel_postnet, mel_linear, stop_token | |
| # 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] | |
| N = 1 | |
| SOS = torch.zeros((N, 1, hp.mel_freq), device=model_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 | |
| 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_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(): | |
| break | |
| else: | |
| # stop_token[:, -1:] is (N, 1) | |
| stop_token_outputs = torch.cat([stop_token_outputs, stop_token[:, -1:]], dim=1) | |
| # 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]) | |
| return mel_to_return, stop_token_outputs | |
| # --- Part 3: Model Loading (from Hugging Face Hub - VERBATIM from your FastAPI code) --- | |
| TTS_MODEL_HUB_ID = "MoHamdyy/transformer-tts-ljspeech" | |
| ASR_HUB_ID = "MoHamdyy/whisper-stt-model" | |
| MARIAN_HUB_ID = "MoHamdyy/marian-ar-en-translation" | |
| 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 | |
| try: | |
| print("Loading TTS model...") | |
| 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.eval() | |
| print("TTS model loaded successfully.") | |
| except Exception as e: print(f"Error loading TTS model: {e}") | |
| 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}") | |
| 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}") | |
| # --- 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 | |
| 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}" | |
| print(msg) | |
| return msg, "", (hp.sr, np.array([]).astype(np.float32)) | |
| # 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) | |
| audio_array_stt = wav.squeeze().cpu().numpy() | |
| print("STT: Extracting features and transcribing...") | |
| inputs = 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, 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}" | |
| # 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:"): | |
| try: | |
| print("TTT: Translating to English...") | |
| batch = mt_tokenizer(arabic_transcript, return_tensors="pt", padding=True).to(DEVICE) | |
| with torch.no_grad(): | |
| 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(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:"): | |
| 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 | |
| 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] | |
| 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"--- 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 | |
| ) | |
| 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 |