Update app.py

app.py

@@ -4,16 +4,16 @@ from transformers import pipeline, WhisperProcessor, WhisperForConditionalGenera
 from gtts import gTTS
 import gradio as gr
 from PIL import Image
-import subprocess
-
-# Import and initialize ZeroGPU
+import logging
+import os
 import spaces
-spaces.init()
 
-print("Using GPU for operations when available")
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# Install flash-attn
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+# Initialize ZeroGPU
+spaces.init()
 
 # Function to safely load pipeline
 @spaces.GPU
@@ -21,7 +21,7 @@ def load_pipeline(model_name, **kwargs):
     try:
         return pipeline(model=model_name, device=0, **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
@@ -32,7 +32,7 @@ def load_whisper():
         model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small").cuda()
         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
@@ -43,10 +43,14 @@ def load_sarvam():
 # Load vision model
 @spaces.GPU
 def load_vision_model():
-    model_id = "microsoft/Phi-3.5-vision-instruct"
-    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, torch_dtype="auto", attn_implementation="flash_attention_2").cuda().eval()
-    processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
-    return model, processor
+    try:
+        model_id = "microsoft/Phi-3.5-vision-instruct"
+        model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, torch_dtype="auto", attn_implementation="flash_attention_2").cuda().eval()
+        processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+        return model, processor
+    except Exception as e:
+        logger.error(f"Error loading vision model: {e}")
+        return None, None
 
 # Process audio input
 @spaces.GPU
@@ -61,7 +65,8 @@ def process_audio_input(audio, whisper_processor, whisper_model):
         transcription = whisper_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
         return transcription
     except Exception as e:
-        return f"Error processing audio: {str(e)}. Please type your message instead."
+        logger.error(f"Error processing audio: {e}")
+        return f"Error processing audio. Please type your message instead."
 
 # Generate response
 def text_to_speech(text, lang='hi'):
@@ -72,10 +77,11 @@ def text_to_speech(text, lang='hi'):
         else:
             tts = gTTS(text=text, lang=lang)
         
-        tts.save("response.mp3")
-        return "response.mp3"
+        output_path = "/tmp/response.mp3"
+        tts.save(output_path)
+        return output_path
     except Exception as e:
-        print(f"Error in text-to-speech: {str(e)}")
+        logger.error(f"Error in text-to-speech: {e}")
         return None
 
 # Detect language (placeholder function, replace with actual implementation)
@@ -93,10 +99,14 @@ def generate_response(transcription, sarvam_pipe):
         response = sarvam_pipe(transcription, max_length=100, num_return_sequences=1)[0]['generated_text']
         return response
     except Exception as e:
-        return f"Error generating response: {str(e)}"
+        logger.error(f"Error generating response: {e}")
+        return f"Error generating response. Please try again."
 
 @spaces.GPU
 def process_image(image, text_input, vision_model, vision_processor):
+    if vision_model is None or vision_processor is None:
+        return "Error: Vision model is not available."
+    
     try:
         prompt = f"<|user|>\n<|image_1|>\n{text_input}<|end|>\n<|assistant|>\n"
         image = Image.fromarray(image).convert("RGB")
@@ -106,7 +116,8 @@ def process_image(image, text_input, vision_model, vision_processor):
         response = vision_processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
         return response
     except Exception as e:
-        return f"Error processing image: {str(e)}"
+        logger.error(f"Error processing image: {e}")
+        return f"Error processing image. Please try again."
 
 @spaces.GPU
 def multimodal_assistant(input_type, audio_input, text_input, image_input):
@@ -131,8 +142,8 @@ def multimodal_assistant(input_type, audio_input, text_input, image_input):
         
         return response, audio_response
     except Exception as e:
-        error_message = f"An error occurred: {str(e)}"
-        return error_message, None
+        logger.error(f"An error occurred in multimodal_assistant: {e}")
+        return f"An error occurred. Please try again.", None
 
 # Custom CSS (you can keep your existing custom CSS here)
 custom_css = """