Update app.py

app.py

@@ -6,25 +6,16 @@ from stitched_model import CombinedModel
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-# Use facebook/mms-1b-all model for ASR and Sunbird/sunbird-mul-en-mbart-merged for translation
+# Load the model
 model = CombinedModel("facebook/mms-1b-all", "Sunbird/sunbird-mul-en-mbart-merged", device=device)
 
-def transcribe(audio_file_mic=None, audio_file_upload=None):
-    if audio_file_mic:
-        audio_file = audio_file_mic
-    elif audio_file_upload:
-        audio_file = audio_file_upload
-    else:
-        return "Please upload an audio file or record one"
-
-    # Load the audio file
+def preprocess_audio(audio_file):
     speech, sample_rate = librosa.load(audio_file, sr=16000, mono=True)
-
-    # Split the audio into 10-second chunks
     chunk_size = 10 * 16000
     chunks = [speech[i:i + chunk_size] for i in range(0, len(speech), chunk_size)]
+    return chunks
 
-    # Process each chunk and concatenate the results
+def transcribe(chunks):
     transcriptions = []
     translations = []
     for chunk in chunks:
@@ -39,9 +30,21 @@ def transcribe(audio_file_mic=None, audio_file_upload=None):
 
     return transcription, translation
 
+def process_audio(audio_file_mic=None, audio_file_upload=None):
+    if audio_file_mic:
+        audio_file = audio_file_mic
+    elif audio_file_upload:
+        audio_file = audio_file_upload
+    else:
+        return "Please upload an audio file or record one"
+
+    chunks = preprocess_audio(audio_file)
+    transcription, translation = transcribe(chunks)
+    return transcription, translation
+
 description = '''Luganda to English Speech Translation'''
 
-iface = gr.Interface(fn=transcribe,
+iface = gr.Interface(fn=process_audio,
                      inputs=[
                          gr.Audio(source="microphone", type="filepath", label="Record Audio"),
                          gr.Audio(source="upload", type="filepath", label="Upload Audio")],