Update app.py

app.py

@@ -3,7 +3,6 @@ import librosa
 from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration, AutoModelForCausalLM, AutoProcessor
 from gtts import gTTS
 import gradio as gr
-import spaces
 from PIL import Image
 import subprocess
 
@@ -12,21 +11,19 @@ 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)
 
-# Function to safely load pipeline within a GPU-decorated function
-@spaces.GPU
+# Function to safely load pipeline
 def load_pipeline(model_name, **kwargs):
     try:
-        device = 0 if torch.cuda.is_available() else "cpu"
+        device = "cuda" 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 Whisper model for speech recognition within a GPU-decorated function
-@spaces.GPU
+# Load Whisper model for speech recognition
 def load_whisper():
     try:
-        device = 0 if torch.cuda.is_available() else "cpu"
+        device = "cuda" 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
@@ -34,21 +31,19 @@ def load_whisper():
         print(f"Error loading Whisper model: {e}")
         return None, None
 
-# Load sarvam-2b for text generation within a GPU-decorated function
-@spaces.GPU
+# Load sarvam-2b for text generation
 def load_sarvam():
     return load_pipeline('sarvamai/sarvam-2b-v0.5')
 
 # 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()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, torch_dtype="auto", attn_implementation="flash_attention_2").to(device).eval()
     processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
     return model, processor
 
-# Process audio input within a GPU-decorated function
-@spaces.GPU
+# Process audio input
 def process_audio_input(audio, whisper_processor, whisper_model):
     if whisper_processor is None or whisper_model is None:
         return "Error: Speech recognition model is not available. Please type your message instead."
@@ -62,8 +57,7 @@ def process_audio_input(audio, whisper_processor, whisper_model):
     except Exception as e:
         return f"Error processing audio: {str(e)}. Please type your message instead."
 
-# Generate response within a GPU-decorated function
-@spaces.GPU
+# Generate response
 def text_to_speech(text, lang='hi'):
     try:
         # Use a better TTS engine for Indic languages
@@ -83,7 +77,6 @@ def detect_language(text):
     # Implement language detection logic here
     return 'en'  # Default to English for now
 
-@spaces.GPU
 def generate_response(transcription, sarvam_pipe):
     if sarvam_pipe is None:
         return "Error: Text generation model is not available."
@@ -95,12 +88,11 @@ def generate_response(transcription, sarvam_pipe):
     except Exception as e:
         return f"Error generating response: {str(e)}"
 
-@spaces.GPU
 def process_image(image, text_input, vision_model, vision_processor):
     try:
         prompt = f"<|user|>\n<|image_1|>\n{text_input}<|end|>\n<|assistant|>\n"
         image = Image.fromarray(image).convert("RGB")
-        inputs = vision_processor(prompt, image, return_tensors="pt").to("cuda:0")
+        inputs = vision_processor(prompt, image, return_tensors="pt").to(vision_model.device)
         generate_ids = vision_model.generate(**inputs, max_new_tokens=1000, eos_token_id=vision_processor.tokenizer.eos_token_id)
         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]
@@ -108,10 +100,9 @@ def process_image(image, text_input, vision_model, vision_processor):
     except Exception as e:
         return f"Error processing image: {str(e)}"
 
-@spaces.GPU
 def multimodal_assistant(input_type, audio_input, text_input, image_input):
     try:
-        # Load models within the GPU-decorated function
+        # Load models
         whisper_processor, whisper_model = load_whisper()
         sarvam_pipe = load_sarvam()
         vision_model, vision_processor = load_vision_model()