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)