Update app.py

app.py

@@ -3,11 +3,12 @@ from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
 from peft import PeftModel
 import spaces
 
+device = "cuda"
 @spaces.GPU
 def greet(image, prompt):
-    base_model = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma-3b-mix-224")
+    base_model = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma-3b-mix-224").to(device)
     processor = AutoProcessor.from_pretrained("google/paligemma-3b-mix-224")
-    model = PeftModel(base_model, "/folders")
+    model = PeftModel(base_model, "/folders").to(device)
     inputs = processor(prompt, raw_image, return_tensors="pt")
     output = model.generate(**inputs, max_new_tokens=20)
     return output