Update app.py

app.py

@@ -1,4 +1,5 @@
 # app.py
+import spaces
 import os
 import gradio as gr
 import torch
@@ -35,7 +36,8 @@ image_model_name = 'resnet50'
 image_embed_dim = 512
 siglip_pretrained_path = "image_encoder.pth" # Path to your pretrained SigLIP model
 
-device = torch.device("cpu") # Force CPU
+#device = torch.device("cpu") # Force CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
 # Load Tokenizer (using a compatible tokenizer)
@@ -71,9 +73,10 @@ image_encoder.eval() # Set to evaluation mode
 #)
 
 base_model_name="microsoft/Phi-3-mini-4k-instruct"
-device = "cuda"
+#device = "cuda"
 
-base_model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.float32, device_map={"": device})
+#base_model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.float32, device_map={"": device})
+base_model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.float32, device_map="auto")
 
 
 # Load and merge
@@ -81,6 +84,8 @@ model = PeftModel.from_pretrained(base_model, peft_model_path, offload_dir='./of
 model = model.merge_and_unload()
 print("phi-3 model loaded sucessfully")  
 # 3. Inference Function
+
+@spaces.GPU
 def predict(image, question):
     """
     Takes an image and a question as input and returns an answer.