Update src/core.py

src/core.py

@@ -12,7 +12,6 @@ import cv2
 
 import numpy as np
 import pandas as pd
-#import streamlit as st
 from PIL import Image
 #from streamlit_drawable_canvas import st_canvas
 
@@ -64,9 +63,9 @@ ENERGY_MASK_CONST = 100000.0              # large energy value for protective ma
 MASK_THRESHOLD = 10                       # minimum pixel intensity for binary mask
 USE_FORWARD_ENERGY = True                 # if True, use forward energy algorithm
 
-device = torch.device("cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model_path = "./assets/big-lama.pt"
-model = torch.jit.load(model_path, map_location="cpu")
+model = torch.jit.load(model_path, map_location=device)
 model = model.to(device)
 model.eval()