Update app.py

app.py

@@ -1,44 +1,37 @@
-import gradio as gr
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForTextToWaveform
+import gradio
 import torch
+import numpy as np
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForTextToWaveform
 
 # Load your pretrained models
 asr_model = Wav2Vec2ForCTC.from_pretrained("Baghdad99/saad-speech-recognition-hausa-audio-to-text")
 asr_processor = Wav2Vec2Processor.from_pretrained("Baghdad99/saad-speech-recognition-hausa-audio-to-text")
-
-# Load the Hausa translation model
 translation_tokenizer = AutoTokenizer.from_pretrained("Baghdad99/saad-hausa-text-to-english-text")
 translation_model = AutoModelForSeq2SeqLM.from_pretrained("Baghdad99/saad-hausa-text-to-english-text", from_tf=True)
-
-# Load the Text-to-Speech model
 tts_tokenizer = AutoTokenizer.from_pretrained("Baghdad99/english_voice_tts")
 tts_model = AutoModelForTextToWaveform.from_pretrained("Baghdad99/english_voice_tts")
 
-def translate_speech(speech):
-    # Extract the audio signal and sample rate
-    audio_signal, sample_rate = speech
-
-    # Convert stereo to mono if necessary
-    if len(audio_signal.shape) > 1:
-        audio_signal = audio_signal.mean(axis=0)
-
-    # Transcribe the speech to text
+# Define the translation and synthesis functions
+def translate(audio_signal):
     inputs = asr_processor(audio_signal, return_tensors="pt", padding=True)
     logits = asr_model(inputs.input_values).logits
     predicted_ids = torch.argmax(logits, dim=-1)
     transcription = asr_processor.decode(predicted_ids[0])
-
-    # Translate the text
     translated = translation_model.generate(**translation_tokenizer(transcription, return_tensors="pt", padding=True))
     translated_text = [translation_tokenizer.decode(t, skip_special_tokens=True) for t in translated]
+    return translated_text
 
-    # Convert the translated text to speech
+def synthesise(translated_text):
     inputs = tts_tokenizer(translated_text, return_tensors='pt')
     audio = tts_model.generate(inputs['input_ids'])
-
     return audio
 
+def translate_speech(audio):
+    translated_text = translate(audio)
+    synthesised_speech = synthesise(translated_text)
+    synthesised_speech = (synthesised_speech.numpy() * max_range).astype(np.int16)
+    return 16000, synthesised_speech
 
 # Define the Gradio interface
-iface = gr.Interface(fn=translate_speech, inputs=gr.inputs.Audio(source="microphone", type="numpy"), outputs="audio")
+iface = gradio.Interface(fn=translate_speech, inputs=gradio.inputs.Audio(source="microphone", type="numpy"), outputs="audio")
 iface.launch()