Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, VitsModel, AutoTokenizer
+import torch
 import numpy as np
 
 # Load the pipeline for speech recognition and translation
@@ -9,7 +10,10 @@ pipe = pipeline(
     tokenizer="Baghdad99/saad-speech-recognition-hausa-audio-to-text"
 )
 translator = pipeline("text2text-generation", model="Baghdad99/saad-hausa-text-to-english-text")
-tts = pipeline("text-to-speech", model="Baghdad99/english_voice_tts")
+
+# Load the VITS model for text-to-speech synthesis
+tts_model = VitsModel.from_pretrained("Baghdad99/english_voice_tts")
+tts_tokenizer = AutoTokenizer.from_pretrained("Baghdad99/english_voice_tts")
 
 # Define the function to translate speech
 def translate_speech(audio):
@@ -39,20 +43,15 @@ def translate_speech(audio):
         print("The translated text does not contain 'generated_token_ids'")
         return
 
-    # Use the text-to-speech pipeline to synthesize the translated text
-    synthesised_speech = tts(translated_text_str)
+    # Use the VITS model to synthesize the translated text
+    tts_inputs = tts_tokenizer(translated_text_str, return_tensors="pt")
+    with torch.no_grad():
+        synthesised_speech = tts_model(**tts_inputs).waveform
     print(f"Synthesised speech: {synthesised_speech}")  # Print the synthesised speech to see what it contains
 
-    # Check if the synthesised speech contains 'audio'
-    if 'audio' in synthesised_speech:
-        synthesised_speech_data = synthesised_speech['audio']
-    else:
-        print("The synthesised speech does not contain 'audio'")
-        return
-
     # Define the max_range variable
     max_range = 1.0  # You can adjust this value based on your requirements
-    synthesised_speech = (synthesised_speech_data * max_range).astype(np.float32)
+    synthesised_speech = (synthesised_speech.numpy() * max_range).astype(np.float32)
 
     return 16000, synthesised_speech