Update app.py

app.py

@@ -9,35 +9,26 @@ print("Using GPU for operations when available")
 
 # Function to safely load pipeline within a GPU-decorated function
 @spaces.GPU
-def indic_language_assistant(input_type, audio_input, text_input, selected_language):
+def load_pipeline(model_name, **kwargs):
     try:
-        # Load models within the GPU-decorated function
-        whisper_processor, whisper_model = load_whisper()
-        sarvam_pipe = load_sarvam()
-
-        if input_type == "audio" and audio_input is not None:
-            transcription = process_audio_input(audio_input, whisper_processor, whisper_model)
-        elif input_type == "text" and text_input:
-            transcription = text_input
-        else:
-            return "Please provide either audio or text input.", "No input provided.", None
+        device = 0 if torch.cuda.is_available() else "cpu"
+        return pipeline(model=model_name, device=device, **kwargs)
+    except Exception as e:
+        print(f"Error loading {model_name} pipeline: {e}")
+        return None
 
-        response = generate_response(transcription, sarvam_pipe)
-        
-        # Convert the selected language to the appropriate language code
-        lang_code = {
-            "Bengali": "bn", "English": "en", "Gujarati": "gu", "Hindi": "hi",
-            "Kannada": "kn", "Malayalam": "ml", "Marathi": "mr", "Oriya": "or",
-            "Punjabi": "pa", "Tamil": "ta", "Telugu": "te"
-        }.get(selected_language, "en")
-        
-        # Use the selected language for text-to-speech
-        audio_response = text_to_speech(response, lang_code)
-        
-        return transcription, response, audio_response
+# Load Whisper model for speech recognition within a GPU-decorated function
+@spaces.GPU
+def load_whisper():
+    try:
+        device = 0 if torch.cuda.is_available() else "cpu"
+        processor = WhisperProcessor.from_pretrained("openai/whisper-small")
+        model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small").to(device)
+        return processor, model
     except Exception as e:
-        error_message = f"An error occurred: {str(e)}"
-        return error_message, error_message, None
+        print(f"Error loading Whisper model: {e}")
+        return None, None
+
 # Load sarvam-2b for text generation within a GPU-decorated function
 @spaces.GPU
 def load_sarvam():
@@ -250,15 +241,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Base().set(
             gr.Markdown("### Indic Assistant")
     
     input_type = gr.Radio(["audio", "text"], label="Input Type", value="audio")
-    
-    with gr.Row():
-        audio_input = gr.Audio(type="filepath", label="Speak (if audio input selected)")
-        language_select = gr.Dropdown(
-            choices=["Bengali", "English", "Gujarati", "Hindi", "Kannada", "Malayalam", "Marathi", "Oriya", "Punjabi", "Tamil", "Telugu"],
-            label="Select Language",
-            value="English"
-        )
-    
+    audio_input = gr.Audio(type="filepath", label="Speak (if audio input selected)")
     text_input = gr.Textbox(label="Type your message (if text input selected)")
     
     submit_btn = gr.Button("Submit")
@@ -269,7 +252,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Base().set(
     
     submit_btn.click(
         fn=indic_language_assistant,
-        inputs=[input_type, audio_input, text_input, language_select],
+        inputs=[input_type, audio_input, text_input],
         outputs=[output_transcription, output_response, output_audio]
     )
     gr.HTML("<footer>Powered by Indic Language AI</footer>")