use whisper-tiny and check processing speeds for low tiers

app.py

@@ -1,56 +1,63 @@
 import gradio as gr
-from transformers import pipeline, SpeechT5Processor, SpeechT5ForTextToSpeech
+from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, AutoModelForSeq2SeqLM, AutoTokenizer
 import torchaudio
 import torch
 from datasets import load_dataset
 import os
 
-# Load ASR and Translation models
-translator = pipeline("translation", model="Helsinki-NLP/opus-mt-en-mul", device=torch.device('cpu'))
-asr = pipeline("automatic-speech-recognition", model="openai/whisper-small", device=torch.device('cpu'))
+# Load lightweight models
+ASR_MODEL = "openai/whisper-tiny"  # Faster ASR model
+TRANSLATION_MODEL = "Helsinki-NLP/opus-mt-en-mul"  # Lightweight translation model
 
-# Load TTS model and processor
+# Load ASR model
+from transformers import pipeline
+asr = pipeline("automatic-speech-recognition", model=ASR_MODEL, device=0 if torch.cuda.is_available() else -1)
+
+# Load translation model and tokenizer
+translator_model = AutoModelForSeq2SeqLM.from_pretrained(TRANSLATION_MODEL)
+translator_tokenizer = AutoTokenizer.from_pretrained(TRANSLATION_MODEL)
+
+# Load TTS processor and model (use float16 for better speed)
 processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
-tts = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
+tts = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts").to(torch.float16)
 
-# Load speaker embeddings from dataset
+# Cache speaker embeddings to avoid reloading every time
 embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
-speaker_embeddings = embeddings_dataset[7306]["xvector"]
-speaker_embeddings = torch.tensor(speaker_embeddings).unsqueeze(0)
+speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0).to(torch.float16)
 
-# Ensure cache directory for output files
+# Ensure output directory exists
 os.makedirs("output", exist_ok=True)
 
-# Function to handle transcription, translation, TTS, and Braille file generation
+# Processing function
 def process_audio(audio, target_language):
+    if not audio:
+        return "Error: No audio file provided.", None, None
+    
     try:
-        if not audio:
-            return "Error: No audio file provided.", None, None
-        
         # Step 1: Transcribe the audio
         result = asr(audio)["text"]
+
         if not result:
             return "Error: Failed to transcribe audio.", None, None
         
         # Step 2: Translate the text
-        translated_text = translator(result, tgt_lang=target_language)
-        if isinstance(translated_text, list):
-            translated_text = translated_text[0].get('translation_text', '')
-        elif isinstance(translated_text, dict):
-            translated_text = translated_text.get('translation_text', '')
+        inputs = translator_tokenizer(result, return_tensors="pt", padding=True)
+        outputs = translator_model.generate(**inputs)
+        translated_text = translator_tokenizer.decode(outputs[0], skip_special_tokens=True)
+
         if not translated_text:
             return "Error: Translation failed.", None, None
 
         # Step 3: Generate speech from translated text
         inputs = processor(text=translated_text, return_tensors="pt")
-        input_features = inputs.input_features
+        input_features = inputs.input_features.to(torch.float16)
 
         with torch.no_grad():
             speech = tts.generate_speech(input_features, speaker_embeddings)
-        
+
         # Save generated speech
         output_audio_path = "output/generated_speech.wav"
-        torchaudio.save(output_audio_path, speech, 24000)
+        torchaudio.save(output_audio_path, speech.cpu(), 24000)
 
         # Step 4: Create Braille-compatible file
         braille_output_path = "output/braille.txt"
@@ -65,7 +72,6 @@ def process_audio(audio, target_language):
 # Define Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Multi-Language Voice Translator")
-    gr.Markdown("Transcribe, translate, and generate speech in multiple languages with accessibility features.")
 
     with gr.Row():
         audio_input = gr.Audio(type="filepath", label="Upload Audio")
@@ -76,20 +82,21 @@ with gr.Blocks() as demo:
         )
 
     with gr.Row():
-        submit_button = gr.Button("Submit")
+        submit_button = gr.Button("Translate & Synthesize")
         clear_button = gr.Button("Clear")
 
     with gr.Row():
         translated_text = gr.Textbox(label="Translated Text")
         generated_speech = gr.Audio(label="Generated Speech", interactive=False)
-        braille_file = gr.File(label="Download Braille-Compatible File")
+        braille_file = gr.File(label="Download Braille File")
 
+    # Link functions to buttons
     submit_button.click(
         fn=process_audio,
         inputs=[audio_input, target_language],
         outputs=[translated_text, generated_speech, braille_file],
     )
-    
+
     clear_button.click(
         fn=lambda: ("", None, None),
         inputs=[],