Update app.py

app.py

@@ -6,51 +6,57 @@ import gradio as gr
 from PIL import Image
 import subprocess
 
+# Import and initialize ZeroGPU
+import spaces
+spaces.init()
+
 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
+@spaces.GPU
 def load_pipeline(model_name, **kwargs):
     try:
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        return pipeline(model=model_name, device=device, **kwargs)
+        return pipeline(model=model_name, device=0, **kwargs)
     except Exception as e:
         print(f"Error loading {model_name} pipeline: {e}")
         return None
 
 # Load Whisper model for speech recognition
+@spaces.GPU
 def load_whisper():
     try:
-        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)
+        model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small").cuda()
         return processor, model
     except Exception as e:
         print(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():
     model_id = "microsoft/Phi-3.5-vision-instruct"
-    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()
+    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
 
 # 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.to(whisper_model.device)
+        input_features = whisper_processor(audio, sampling_rate=sr, return_tensors="pt").input_features.cuda()
         predicted_ids = whisper_model.generate(input_features)
         transcription = whisper_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
         return transcription
@@ -77,6 +83,7 @@ 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."
@@ -88,11 +95,12 @@ 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(vision_model.device)
+        inputs = vision_processor(prompt, image, return_tensors="pt").to("cuda")
         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]
@@ -100,6 +108,7 @@ 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