Add model checkpoints

app.py

@@ -75,7 +75,11 @@ specialist_model, specialist_processor = build_specialist_model()
 
 
 def load_specialist_model(model_ckpt):
-    specialist_model.load_state_dict(torch.load(model_ckpt, map_location=torch.device("cpu"), weights_only=True))
+    try:
+        specialist_model.load_state_dict(torch.load(model_ckpt, map_location=torch.device("cpu"), weights_only=True))
+        print(f"Loaded {model_ckpt}")
+    except:
+        print(f"Failed to load {model_ckpt}")
 
 
 def get_feature_dict(batch_size, device, active_dataset: ActiveDataset):
@@ -169,7 +173,11 @@ def build_input_ui():
 def build_parameters_ui():
     with gr.Accordion() as blk:
         budget_input = gr.Number(config.budget, label="Budget")
-        model_ckpt_input = gr.Text(config.model_ckpt, label="Specialist Model Checkpoint")
+
+        with gr.Row():
+            model_ckpt_file = gr.File(visible=False)
+            model_ckpt_input = gr.UploadButton(label="Upload model checkpoint")
+
         device_input = gr.Dropdown(choices=["cuda", "cpu"], value="cpu", label="Device", interactive=True)
         batch_size_input = gr.Number(config.batch_size, label="Batch Size")
         foundation_model_weight_input = gr.Number(config.loaded_feature_weight, label="foundation_model_weight")
@@ -180,10 +188,11 @@ def build_parameters_ui():
 
         budget_input.change(budget_input_change, budget_input, None)
 
-        def model_ckpt_input_change(x):
+        def model_ckpt_input_upload(x):
             config.model_ckpt = x
+            return gr.File(value=x, visible=True)
 
-        model_ckpt_input.change(model_ckpt_input_change, model_ckpt_input, None)
+        model_ckpt_input.upload(model_ckpt_input_upload, model_ckpt_input, model_ckpt_file)
 
         def device_input_change(x):
             config.device = torch.device(x)