Update app.py

app.py

@@ -21,8 +21,12 @@ def load_model():
         "google/paligemma2-28b-pt-896", use_auth_token=token
     )
     model = AutoModelForImageTextToText.from_pretrained(
-        "google/paligemma2-28b-pt-896", use_auth_token=token
+        "google/paligemma2-28b-pt-896", use_auth_token=token, torch_dtype=torch.float16
     )
+    
+    # Move model to GPU if available
+    if torch.cuda.is_available():
+        model = model.to("cuda")
 
     return processor, model
 
@@ -33,7 +37,7 @@ def process_image(image):
     processor, model = load_model()
 
     # Preprocess the image
-    inputs = processor(images=image, return_tensors="pt")
+    inputs = processor(images=image, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu", dtype=torch.float16)
 
     # Generate predictions
     with torch.no_grad():