Cache loaded model

app.py

@@ -12,11 +12,24 @@ from isegm.inference import clicker as ck
 from isegm.inference import utils
 from isegm.inference.predictors import get_predictor
 
-# Model Path
-prefix = "https://huggingface.co/curt-park/interactive-segmentation/resolve/main"
-models = {
-    "RITM": "ritm_coco_lvis_h18_itermask.pth",
-}
+@st.cache_data
+def load_model(model_path, device):
+    model = utils.load_is_model(model_path, device, cpu_dist_maps=True)
+    predictor_params = {"brs_mode": "NoBRS"}
+    predictor = get_predictor(model, device=device, **predictor_params)
+    return predictor
+
+
+# Objects in the global scope
+url_prefix = "https://huggingface.co/curt-park/interactive-segmentation/resolve/main"
+models = {"RITM": "ritm_coco_lvis_h18_itermask.pth"}
+clicker = ck.Clicker()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+pos_color, neg_color = "#3498DB", "#C70039"
+canvas_height, canvas_width = 600, 600
+err_x, err_y = 5.5, 1.0
+predictor = None
+image = None
 
 # Items in the sidebar.
 model = st.sidebar.selectbox("Select a Model:", tuple(models.keys()))
@@ -25,24 +38,16 @@ marking_type = st.sidebar.radio("Marking Type:", ("positive", "negative"))
 image_path = st.sidebar.file_uploader("Background Image:", type=["png", "jpg", "jpeg"])
 
 # Objects for prediction.
-clicker = ck.Clicker()
-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]):
-        _ = wget.download(f"{prefix}/{models[model]}")
+        _ = wget.download(f"{url_prefix}/{models[model]}")
 
 with st.spinner("Wait for loading a model..."):
-    model = utils.load_is_model(models[model], device, cpu_dist_maps=True)
-    predictor_params = {"brs_mode": "NoBRS"}
-    predictor = get_predictor(model, device=device, **predictor_params)
+    predictor = load_model(models[model], device)
 
 # Create a canvas component.
-image = None
 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(
@@ -66,7 +71,6 @@ if canvas_result.json_data and canvas_result.json_data["objects"] and image:
     image_width, image_height = image.size
     ratio_h, ratio_w = image_height / canvas_height, image_width / canvas_width
 
-    err_x, err_y = 5.5, 1.0
     pos_clicks, neg_clicks = [], []
     for click in objects:
         x, y = (click["left"] + err_x) * ratio_w, (click["top"] + err_y) * ratio_h