Update app.py

app.py

@@ -13,9 +13,10 @@ print("Using GPU for operations when available")
 def load_pipeline(model_name, **kwargs):
     try:
         device = 0 if torch.cuda.is_available() else "cpu"
+        logger.info(f"Loading {model_name} on device: {device}")
         return pipeline(model=model_name, device=device, **kwargs)
     except Exception as e:
-        print(f"Error loading {model_name} pipeline: {e}")
+        logger.error(f"Error loading {model_name} pipeline: {e}")
         return None
 
 # Load Whisper model for speech recognition within a GPU-decorated function
@@ -23,18 +24,31 @@ def load_pipeline(model_name, **kwargs):
 def load_whisper():
     try:
         device = 0 if torch.cuda.is_available() else "cpu"
+        logger.info(f"Loading Whisper model on device: {device}")
         processor = WhisperProcessor.from_pretrained("openai/whisper-small")
         model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small").to(device)
         return processor, model
     except Exception as e:
-        print(f"Error loading Whisper model: {e}")
+        logger.error(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():
+    logger.info("Loading sarvam-2b model")
     return load_pipeline('sarvamai/sarvam-2b-v0.5')
 
+
+# Global variables for models
+whisper_processor, whisper_model = load_whisper()
+sarvam_pipe = load_sarvam()
+
+# Check if models are loaded
+if whisper_processor is None or whisper_model is None:
+    logger.error("Whisper model failed to load")
+if sarvam_pipe is None:
+    logger.error("Sarvam model failed to load")
+
 # Process audio input within a GPU-decorated function
 @spaces.GPU
 def process_audio_input(audio, whisper_processor, whisper_model):
@@ -106,16 +120,17 @@ def detect_language(text):
 @spaces.GPU
 def indic_language_assistant(input_type, audio_input, text_input):
     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:
+            if whisper_processor is None or whisper_model is None:
+                return "Error: Speech recognition model is not available.", "", 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
+            return "Please provide either audio or text input.", "", None
+
+        if sarvam_pipe is None:
+            return transcription, "Error: Text generation model is not available.", None
 
         response = generate_response(transcription, sarvam_pipe)
         lang = detect_language(response)
@@ -123,9 +138,8 @@ def indic_language_assistant(input_type, audio_input, text_input):
         
         return transcription, response, audio_response
     except Exception as e:
-        error_message = f"An error occurred: {str(e)}"
-        return error_message, error_message, None
-
+        logger.error(f"An error occurred in indic_language_assistant: {str(e)}")
+        return str(e), "An error occurred while processing your request.", None
 
 
 # Updated Custom CSS