Load the speaker embeddings directly from the dataset

app.py

@@ -1,12 +1,22 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, SpeechT5Processor, SpeechT5ForTextToSpeech
+import torchaudio
 import torch
+from datasets import load_dataset
 import os
 
-# Initialize models
+# Load ASR and Translation models
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-en-mul", device=torch.device('cpu'))
 asr = pipeline("automatic-speech-recognition", model="openai/whisper-small", device=torch.device('cpu'))
-tts = pipeline("text-to-speech", model="microsoft/speecht5_tts", device=torch.device('cpu'))
+
+# Load TTS model and processor
+processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
+tts = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
+
+# Load speaker embeddings from dataset
+embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
+speaker_embeddings = embeddings_dataset[7306]["xvector"]  # Example speaker embedding
+speaker_embeddings = torch.tensor(speaker_embeddings).unsqueeze(0)  # Reshape for the model
 
 # Ensure cache directory for output files
 os.makedirs("output", exist_ok=True)
@@ -25,9 +35,14 @@ def process_audio(audio, target_language):
             translated_text = translated_text.get('translation_text', '')
 
         # Step 3: Generate speech from translated text
-        speech = tts(translated_text)
+        inputs = processor(translated_text, return_tensors="pt")
+
+        with torch.no_grad():
+            speech = tts.generate_speech(inputs["input_ids"], speaker_embeddings)
+        
+        # Save generated speech
         output_audio_path = "output/generated_speech.wav"
-        speech["audio"].save(output_audio_path)
+        torchaudio.save(output_audio_path, speech, 24000)
 
         # Step 4: Create Braille-compatible file
         braille_output_path = "output/braille.txt"
@@ -46,7 +61,11 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         audio_input = gr.Audio(type="filepath", label="Upload Audio")
-        target_language = gr.Dropdown(choices=["en", "hi", "kn", "ta", "te", "es", "de", "fr", "hu"], value="en", label="Target Language")
+        target_language = gr.Dropdown(
+            choices=["en", "hi", "kn", "ta", "te", "es", "de", "fr", "hu"],
+            value="en",
+            label="Target Language"
+        )
 
     with gr.Row():
         submit_button = gr.Button("Submit")