Update app.py

app.py

@@ -1,72 +1,106 @@
+# Import spaces first to avoid CUDA initialization issues
+import spaces
+
+# Then import other libraries
 import torch
 import librosa
-from transformers import AutoModelForCausalLM, AutoProcessor, pipeline, WhisperProcessor, WhisperForConditionalGeneration
+from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration, AutoModelForCausalLM, AutoProcessor
 from gtts import gTTS
 import gradio as gr
-import spaces
 from PIL import Image
 import os
 from langdetect import detect
 import subprocess
 
+print("Using GPU for operations when available")
+
 # Install flash-attn
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-print("Loading models...")
+# Function to safely load pipeline within a GPU-decorated function
+@spaces.GPU
+def load_pipeline(model_name, **kwargs):
+    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
 
-# Vision model
-vision_model_id = "microsoft/Phi-3.5-vision-instruct"
-vision_model = AutoModelForCausalLM.from_pretrained(
-    vision_model_id,
-    trust_remote_code=True,
-    torch_dtype=torch.float16,
-    use_flash_attention_2=False
-)
-vision_processor = AutoProcessor.from_pretrained(vision_model_id, trust_remote_code=True, num_crops=16)
+# 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:
+        print(f"Error loading Whisper model: {e}")
+        return None, None
 
-# Whisper model
-whisper_model_id = "openai/whisper-small"
-whisper_processor = WhisperProcessor.from_pretrained(whisper_model_id)
-whisper_model = WhisperForConditionalGeneration.from_pretrained(whisper_model_id)
+# Load vision model within a GPU-decorated function
+@spaces.GPU
+def load_vision_model():
+    try:
+        model_id = "microsoft/Phi-3.5-vision-instruct"
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id, trust_remote_code=True, torch_dtype=torch.float16, use_flash_attention_2=False
+        )
+        processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True, num_crops=16)
+        return model, processor
+    except Exception as e:
+        print(f"Error loading vision model: {e}")
+        return None, None
 
-# Sarvam model
-sarvam_pipe = pipeline('sarvamai/sarvam-2b-v0.5')
+# Load sarvam-2b for text generation within a GPU-decorated function
+@spaces.GPU
+def load_sarvam():
+    return load_pipeline('sarvamai/sarvam-2b-v0.5')
 
-print("All models loaded successfully")
+# Load all models
+whisper_processor, whisper_model = load_whisper()
+vision_model, vision_processor = load_vision_model()
+sarvam_pipe = load_sarvam()
 
 @spaces.GPU
 def process_audio_input(audio):
+    if whisper_processor is None or whisper_model is None:
+        return "Error: Speech recognition model is not available. Please type your message instead."
+    
     try:
-        whisper_model.to('cuda')
         audio, sr = librosa.load(audio, sr=16000)
-        input_features = whisper_processor(audio, sampling_rate=sr, return_tensors="pt").input_features.to('cuda')
+        input_features = whisper_processor(audio, sampling_rate=sr, return_tensors="pt").input_features.to(whisper_model.device)
         predicted_ids = whisper_model.generate(input_features)
         transcription = whisper_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
-        whisper_model.to('cpu')
         return transcription
     except Exception as e:
         return f"Error processing audio: {str(e)}. Please type your message instead."
 
 @spaces.GPU
 def process_image_input(image, text_prompt):
+    if vision_model is None or vision_processor is None:
+        return "Error: Vision model is not available."
+    
     try:
-        vision_model.to('cuda')
         messages = [
             {"role": "user", "content": f"{text_prompt}\n<|image_1|>"},
         ]
         prompt = vision_processor.tokenizer.apply_chat_template(
             messages, tokenize=False, add_generation_prompt=True
         )
-        inputs = vision_processor(prompt, image, return_tensors="pt").to("cuda")
+        inputs = vision_processor(prompt, image, return_tensors="pt").to(vision_model.device)
         generate_ids = vision_model.generate(**inputs, max_new_tokens=1000, temperature=0.2, do_sample=True)
         generate_ids = generate_ids[:, inputs['input_ids'].shape[1]:]
         response = vision_processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-        vision_model.to('cpu')
         return response
     except Exception as e:
         return f"Error processing image: {str(e)}"
 
 def generate_response(transcription):
+    if sarvam_pipe is None:
+        return "Error: Text generation model is not available."
+    
     try:
         response = sarvam_pipe(transcription, max_length=100, num_return_sequences=1)[0]['generated_text']
         return response