Update app.py

app.py

@@ -6,46 +6,42 @@ import gradio as gr
 from PIL import Image
 import logging
 import os
-import spaces
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Initialize ZeroGPU
-spaces.init()
+# Check for GPU availability
+device = "cuda" if torch.cuda.is_available() else "cpu"
+logger.info(f"Using device: {device}")
 
 # Function to safely load pipeline
-@spaces.GPU
 def load_pipeline(model_name, **kwargs):
     try:
-        return pipeline(model=model_name, device=0, **kwargs)
+        return pipeline(model=model_name, device=device, **kwargs)
     except Exception as e:
         logger.error(f"Error loading {model_name} pipeline: {e}")
         return None
 
 # Load Whisper model for speech recognition
-@spaces.GPU
 def load_whisper():
     try:
         processor = WhisperProcessor.from_pretrained("openai/whisper-small")
-        model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small").cuda()
+        model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small").to(device)
         return processor, model
     except Exception as e:
         logger.error(f"Error loading Whisper model: {e}")
         return None, None
 
 # Load sarvam-2b for text generation
-@spaces.GPU
 def load_sarvam():
     return load_pipeline('sarvamai/sarvam-2b-v0.5')
 
 # Load vision model
-@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="auto", attn_implementation="flash_attention_2").cuda().eval()
+        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
     except Exception as e:
@@ -53,14 +49,13 @@ def load_vision_model():
         return None, None
 
 # Process audio input
-@spaces.GPU
 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."
     
     try:
         audio, sr = librosa.load(audio, sr=16000)
-        input_features = whisper_processor(audio, sampling_rate=sr, return_tensors="pt").input_features.cuda()
+        input_features = whisper_processor(audio, sampling_rate=sr, return_tensors="pt").input_features.to(device)
         predicted_ids = whisper_model.generate(input_features)
         transcription = whisper_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
         return transcription
@@ -89,7 +84,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."
@@ -102,7 +96,6 @@ def generate_response(transcription, sarvam_pipe):
         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."
@@ -110,7 +103,7 @@ 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")
+        inputs = vision_processor(prompt, image, return_tensors="pt").to(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]
@@ -119,7 +112,6 @@ def process_image(image, text_input, vision_model, vision_processor):
         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):
     try:
         # Load models