Make gpu compatible

app.py

@@ -26,7 +26,7 @@ image_path = st.sidebar.file_uploader("Background Image:", type=["png", "jpg", "
 
 # Objects for prediction.
 clicker = ck.Clicker()
-device = torch.device("cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 predictor = None
 with st.spinner("Wait for downloading a model..."):
     if not os.path.exists(models[model]):
@@ -43,6 +43,7 @@ if image_path:
     image = Image.open(image_path).convert("RGB")
 canvas_height, canvas_width = 600, 600
 pos_color, neg_color = "#3498DB", "#C70039"
+
 st.title("Canvas:")
 canvas_result = st_canvas(
         fill_color="rgba(255, 165, 0, 0.3)",  # Fixed fill color with some opacity
@@ -75,11 +76,15 @@ if canvas_result.json_data and canvas_result.json_data["objects"] and image:
         click = ck.Click(is_positive=is_positive, coords=(y, x))
         clicker.add_click(click)
 
-    # prediction.
+    # Run prediction.
     pred = None
     predictor.set_input_image(np.array(image))
+    init_mask = torch.zeros((1, 1, image_height, image_width), device=device)
+
     with st.spinner("Wait for prediction..."):
-        pred = predictor.get_prediction(clicker, prev_mask=None)
+        pred = predictor.get_prediction(clicker, prev_mask=init_mask)
     pred = cv2.resize(pred, dsize=(canvas_height, canvas_width), interpolation=cv2.INTER_CUBIC)
     pred = np.where(pred > threshold, 1.0, 0)
+
+    # Show the prediction result.
     st.image(pred, caption="")