app.py

import gradio as gr
import torch
import os
import librosa
from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5ForSpeechToText
from peft import PeftModel, PeftConfig
from datasets import load_dataset, Audio
import numpy as np
from speechbrain.inference import EncoderClassifier

# Initialize Global Device
device = "cuda" if torch.cuda.is_available() else "cpu"

print("Loading Processor...")
processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts") 

print("Loading Base Models...")
base_tts_model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts").to(device)
base_stt_model = SpeechT5ForSpeechToText.from_pretrained("microsoft/speecht5_asr").to(device)

print("Injecting Custom PEFT LoRA Adapters...")
try:
    # We dynamically strip out the "SEQ_2_SEQ_LM" task_type locally! 
    # If we don't, PEFT falsely assumes SpeechT5 is a generic Text Language Model and crashes looking for 'prepare_inputs_for_generation'.
    tts_repo = "Solo448/SpeechT5-Unified-TTS-PEFT"
    tts_config = PeftConfig.from_pretrained(tts_repo)
    tts_config.task_type = None 
    tts_model = PeftModel.from_pretrained(base_tts_model, tts_repo, config=tts_config).to(device)
    
    stt_repo = "Solo448/SpeechT5-Unified-STT-PEFT"
    stt_config = PeftConfig.from_pretrained(stt_repo)
    stt_config.task_type = None
    stt_model = PeftModel.from_pretrained(base_stt_model, stt_repo, config=stt_config).to(device)
except Exception as e:
    print(f"Warning! Failed to load custom adapters with Config bypass: {e}")

print("Loading Vocoder...")
vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan").to(device)

print("Loading Speaker Model...")
speaker_model = EncoderClassifier.from_hparams(
    source="speechbrain/spkrec-xvect-voxceleb",
    run_opts={"device": device},
    savedir=os.path.join("/tmp", "speechbrain/spkrec-xvect-voxceleb")
)

# Speed optimization: Load via streaming so the UI boots up in seconds instead of permanently hanging while downloading Mozilla gigabytes
print("Preparing Speaker Embedding...")
try:
    dataset = load_dataset("mozilla-foundation/common_voice_17_0", "hi", split="validated", trust_remote_code=True, streaming=True)
    sample = next(iter(dataset))
    audio_array = librosa.resample(sample['audio']['array'], orig_sr=sample['audio']['sampling_rate'], target_sr=16000)
    
    def create_speaker_embedding(waveform):
        with torch.no_grad():
            speaker_embeddings = speaker_model.encode_batch(torch.tensor(waveform))
            speaker_embeddings = torch.nn.functional.normalize(speaker_embeddings, dim=2)
            speaker_embeddings = speaker_embeddings.squeeze().cpu().numpy()
        return speaker_embeddings
        
    speaker_embedding = torch.tensor(create_speaker_embedding(audio_array)).unsqueeze(0).to(device)
    print("Speaker Embedding Ready!")
except Exception as e:
    print(f"Error fetching speaker sample dynamically: {e}")
    speaker_embedding = torch.randn(1, 512).to(device) # Metal static fallback

# Unified Transliteration Dictionary
replacements = [
    # Bengali
    ("অ", "a"), ("আ", "aa"), ("ই", "i"), ("ঈ", "ee"), ("উ", "u"), ("ঊ", "oo"), ("ঋ", "ri"), ("এ", "e"), ("ঐ", "oi"), ("ও", "o"), ("ঔ", "ou"),
    ("ক", "k"), ("খ", "kh"), ("গ", "g"), ("ঘ", "gh"), ("ঙ", "ng"), ("চ", "ch"), ("ছ", "chh"), ("জ", "j"), ("ঝ", "jh"), ("ঞ", "nj"),
    ("ট", "t"), ("ঠ", "th"), ("ড", "d"), ("ঢ", "dh"), ("ণ", "nr"), ("ত", "t"), ("थ", "th"), ("দ", "d"), ("ध", "dh"), ("न", "n"), ("प", "p"),
    ("ফ", "ph"), ("ব", "b"), ("ভ", "bh"), ("ম", "m"), ("য", "ya"), ("র", "r"), ("ল", "l"), ("শ", "sha"), ("ষ", "sh"), ("স", "s"), ("হ", "ha"),
    ("ড়", "rh"), ("ঢ়", "rh"), ("য়", "y"), ("ৎ", "t"), ("ঃ", "h"), ("ঁ", "n"), ("़", ""), ("া", "a"), ("ি", "i"), ("ী", "ii"), ("ু", "u"), ("ূ", "uu"),
    ("ৃ", "r"), ("ে", "e"), ("ৈ", "oi"), ("ো", "o"), ("ৌ", "ou"), ("্", ""), ("ৎ", "t"), ("ৗ", "ou"), ("ড়", "r"), ("ঢ়", "r"), ("য়", "y"), ("ৰ", "r"), ("৵", "lee"), ("ং", "ng"), ("১", "1"), ("২", "2"), ("৩", "3"), ("৪", "4"), ("৫", "5"), ("৬", "6"), ("৭", "7"), ("৮", "8"), ("৯", "9"), ("০", "0"),
    # Hindi
    ("अ", "a"), ("आ", "aa"), ("इ", "i"), ("ई", "ee"), ("उ", "u"), ("ऋ", "ri"), ("ए", "ae"), ("ऐ", "ai"), ("ऑ", "au"), ("ओ", "o"), ("औ", "au"),
    ("क", "k"), ("ख", "kh"), ("ग", "g"), ("घ", "gh"), ("च", "ch"), ("छ", "chh"), ("ज", "j"), ("झ", "jh"), ("ञ", "gna"), ("ट", "t"), ("ठ", "th"),
    ("ड", "d"), ("ढ", "dh"), ("ण", "nr"), ("त", "t"), ("थ", "th"), ("द", "d"), ("ध", "dh"), ("न", "n"), ("प", "p"), ("फ", "ph"), ("ब", "b"),
    ("भ", "bh"), ("म", "m"), ("य", "ya"), ("र", "r"), ("ल", "l"), ("व", "w"), ("श", "sha"), ("ष", "sh"), ("स", "s"), ("ह", "ha"),
    ("़", "ng"), ("ऽ", ""), ("ा", "a"), ("ि", "i"), ("ी", "ee"), ("ु", "u"), ("ॅ", "n"), ("े", "e"), ("ै", "oi"), ("ो", "o"), ("ौ", "ou"),
    ("ॅ", "n"), ("ॉ", "r"), ("ू", "uh"), ("ृ", "ri"), ("ं", "n"), ("क़", "q"), ("ज़", "z"), ("ड़", "r"), ("ढ़", "rh"), ("फ़", "f"), ("|", ".")
]

def clean_text(text):
    for src, dst in replacements:
        text = text.replace(src, dst)
    return text.lower() # ASR natively relies on lowercased targets

def text_to_speech_fn(text):
    if not text.strip():
        return None
    cleaned_txt = clean_text(text)
    inputs = processor(text=cleaned_txt, return_tensors="pt").to(device)
    speech = tts_model.generate_speech(inputs["input_ids"], speaker_embedding, vocoder=vocoder)
    return (16000, speech.cpu().numpy())

def speech_to_text_fn(audio_path):
    if not audio_path:
        return "Please provide an audio sample."
        
    audio_array, sr = librosa.load(audio_path, sr=16000)
    
    inputs = processor(audio=audio_array, sampling_rate=16000, return_tensors="pt").to(device)
    predicted_ids = stt_model.generate(**inputs, max_length=150)
    
    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
    return transcription

with gr.Blocks(title="Unified Indic SpeechT5") as demo:
    gr.Markdown("# 🗣️ Unified Multi-Lingual SpeechT5 (Hindi & Bengali)")
    gr.Markdown("This interface dynamically taps into Custom PEFT LoRA Adapters for both Text-to-Speech and Speech-To-Text across 2 massive linguistic domains natively without architectural compromises!")
    
    with gr.Tabs():
        # Text-to-Speech Tab
        with gr.TabItem("Text to Speech (TTS)"):
            with gr.Row():
                tts_input = gr.Textbox(label="Enter Hindi or Bengali Text", lines=3, placeholder="मैं एक कृत्रिम बुद्धिमत्ता हूँ...")
            with gr.Row():
                tts_btn = gr.Button("Generate Speech", variant="primary")
            with gr.Row():
                tts_output = gr.Audio(label="Synthesized Audio output", autoplay=False)
            
            tts_btn.click(fn=text_to_speech_fn, inputs=tts_input, outputs=tts_output)

        # Speech-to-Text Tab
        with gr.TabItem("Speech to Text (ASR)"):
            gr.Markdown("*(Note: Due to the tokenizer hack natively applied during fine-tuning, the acoustic transcription yields mathematically precise Romanized / Latin phonetic representations of your Hindi/Bengali speech.)*")
            with gr.Row():
                stt_input = gr.Audio(label="Record or Upload Spoken Audio", type="filepath")
            with gr.Row():
                stt_btn = gr.Button("Transcribe Audio", variant="primary")
            with gr.Row():
                stt_output = gr.Textbox(label="Transcribed Romanized Text", interactive=False)
                
            stt_btn.click(fn=speech_to_text_fn, inputs=stt_input, outputs=stt_output)

if __name__ == "__main__":
    demo.launch(share=True)