Update app.py

app.py

@@ -2,41 +2,46 @@ import gradio as gr
 import os
 import torch
 import soundfile as sf
-import torchaudio
-from transformers import VitsModel, AutoTokenizer
-from speechbrain.pretrained import EncoderDecoderASR
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, VitsModel, AutoTokenizer
 
 # Ensure the output directory exists
 os.makedirs("output_audio", exist_ok=True)
 
-# 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")
+# Load the models and processors
+asr_processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-large-960h")
+asr_model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-large-960h")
+tts_model = VitsModel.from_pretrained("facebook/mms-tts-eng")
+tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng")
 
 def speech_to_text(input_audio):
-    sig, sr = torchaudio.load(input_audio)
-    if sr != 16000:
-        sig = torchaudio.transforms.Resample(orig_freq=sr, new_freq=16000)(sig)
-    transcription = asr_model.transcribe_file(input_audio)
+    # Load and preprocess the audio
+    waveform, sr = sf.read(input_audio)
+    input_values = asr_processor(waveform, sampling_rate=sr, return_tensors="pt").input_values
+    
+    # Perform speech recognition
+    with torch.no_grad():
+        logits = asr_model(input_values).logits
+    predicted_ids = torch.argmax(logits, dim=-1)
+    
+    # Decode the predicted IDs to text
+    transcription = asr_processor.batch_decode(predicted_ids)[0]
     return transcription
 
 def text_to_speech(text):
-    inputs = tokenizer(text, return_tensors="pt")
+    # Tokenize text and generate waveform
+    inputs = tts_tokenizer(text, return_tensors="pt")
     with torch.no_grad():
-        output = model(**inputs).waveform
+        output = tts_model(**inputs).waveform
     waveform = output.numpy()
 
+    # Define output path and save waveform as audio file
     output_path = "output_audio/text_to_speech.wav"
     sf.write(output_path, waveform.squeeze(), 22050)
 
     return output_path
 
 def speech_to_speech(input_audio, target_text):
-    # Use speech_to_text to transcribe, then synthesize speech from the transcription
-    transcription = speech_to_text(input_audio)
+    # Synthesize speech directly from target text without transcribing the input audio
     return text_to_speech(target_text)
 
 iface = gr.Interface(
@@ -55,11 +60,10 @@ iface = gr.Interface(
         "Text to Speech": gr.outputs.Audio(type="file", label="Synthesized Speech"),
         "Speech to Speech": gr.outputs.Audio(type="file", label="Synthesized Speech")
     },
-    title="Speech Processing App",
-    description="This app uses SpeechBrain for speech to text and Facebook's MMS for text to speech.",
+    title="Speech Processing Application",
+    description="This app uses Facebook's Wav2Vec 2.0 for speech-to-text and VITS for text-to-speech.",
     layout="vertical"
 )
 
 if __name__ == "__main__":
     iface.launch()
-