Update app.py

app.py

@@ -3,16 +3,15 @@ import os
 import torch
 import soundfile as sf
 import torchaudio
-from scipy.io.wavfile import write
-from transformers import VitsProcessor, VitsForConditionalGeneration
-from speechbrain.pretrained import EncoderClassifier, EncoderDecoderASR
+from transformers import VitsModel, AutoTokenizer
+from speechbrain.pretrained import EncoderDecoderASR
 
 # Ensure the output directory exists
 os.makedirs("output_audio", exist_ok=True)
 
-# Load the Facebook MMS TTS model and processor
-tts_processor = VitsProcessor.from_pretrained("facebook/mms-tts-eng")
-tts_model = VitsForConditionalGeneration.from_pretrained("facebook/mms-tts-eng")
+# Load the Facebook MMS TTS model and tokenizer
+model = VitsModel.from_pretrained("facebook/mms-tts-eng")
+tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng")
 
 # SpeechBrain ASR Model for Speech to Text
 asr_model = EncoderDecoderASR.from_hparams(source="speechbrain/asr-conformer-transformerlm-librispeech", savedir="models/asr")
@@ -25,20 +24,19 @@ def speech_to_text(input_audio):
     return transcription
 
 def text_to_speech(text):
-    inputs = tts_processor(text, return_tensors="pt")
+    inputs = tokenizer(text, return_tensors="pt")
     with torch.no_grad():
-        generated = tts_model.generate(**inputs)
-        waveform = generated.audio.squeeze().cpu().numpy()
+        output = model(**inputs).waveform
+    waveform = output.numpy()
 
     output_path = "output_audio/text_to_speech.wav"
-    sf.write(output_path, waveform, 22050)
+    sf.write(output_path, waveform.squeeze(), 22050)
+
     return output_path
 
 def speech_to_speech(input_audio, target_text):
-    # Speech to Text
+    # Use speech_to_text to transcribe, then synthesize speech from the transcription
     transcription = speech_to_text(input_audio)
-    
-    # Text to Speech with Facebook MMS TTS
     return text_to_speech(target_text)
 
 iface = gr.Interface(
@@ -64,3 +62,4 @@ iface = gr.Interface(
 
 if __name__ == "__main__":
     iface.launch()
+