Update app.py

app.py

@@ -9,25 +9,28 @@ print("Using GPU for operations when available")
 
 # Function to safely load pipeline within a GPU-decorated function
 @spaces.GPU
-def load_pipeline(model_name, **kwargs):
+def indic_language_assistant(input_type, audio_input, text_input, selected_language):
     try:
-        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
+        # Load models within the GPU-decorated function
+        whisper_processor, whisper_model = load_whisper()
+        sarvam_pipe = load_sarvam()
 
-# 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
+        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
+
+        response = generate_response(transcription, sarvam_pipe)
+        
+        # Use the selected language for text-to-speech
+        audio_response = text_to_speech(response, selected_language)
+        
+        return transcription, response, audio_response
     except Exception as e:
-        print(f"Error loading Whisper model: {e}")
-        return None, None
+        error_message = f"An error occurred: {str(e)}"
+        return error_message, error_message, None
 
 # Load sarvam-2b for text generation within a GPU-decorated function
 @spaces.GPU
@@ -241,7 +244,15 @@ 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")
-    audio_input = gr.Audio(type="filepath", label="Speak (if audio input selected)")
+    
+    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"
+        )
+    
     text_input = gr.Textbox(label="Type your message (if text input selected)")
     
     submit_btn = gr.Button("Submit")
@@ -252,7 +263,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],
+        inputs=[input_type, audio_input, text_input, language_select],
         outputs=[output_transcription, output_response, output_audio]
     )
     gr.HTML("<footer>Powered by Indic Language AI</footer>")