Update app.py

app.py

@@ -150,7 +150,7 @@ def build_model_and_transforms(args):
     build_func = MODULE_BUILD_FUNCS.get(args.modelname)
     model, _, _ = build_func(args)
 
-    model.to(device)
+    #model.to(device)
 
     checkpoint = torch.load(args.pretrain_model_path, map_location="cpu")["model"]
     model.load_state_dict(checkpoint, strict=False)
@@ -223,6 +223,7 @@ def get_ind_to_filter(text, word_ids, keywords):
 
 @spaces.GPU
 def count(image, text, prompts, state, device):
+    model.to(device)
     print("state: " + str(state))
     keywords = "" # do not handle this for now
     # Handle no prompt case.
@@ -308,10 +309,12 @@ def count(image, text, prompts, state, device):
             out_label += " " + str(exemplars[0].size()[0]) + " visual exemplars."
     else:
         out_label = "Nothing specified to detect."
+    model.cpu()
     return (gr.Image(output_img, visible=True, label=out_label, show_label=True), gr.Number(label="Predicted Count", visible=True, value=boxes.shape[0]), new_submit_btn, gr.Tab(visible=True), step_3, state)
 
 @spaces.GPU
 def count_main(image, text, prompts, device):
+    model.to(device)
     keywords = "" # do not handle this for now
     # Handle no prompt case.
     if prompts is None:
@@ -378,6 +381,7 @@ def count_main(image, text, prompts, device):
             out_label += " " + str(exemplars[0].size()[0]) + " visual exemplars."
     else:
         out_label = "Nothing specified to detect."
+    model.cpu()
     return (gr.Image(output_img, visible=True, label=out_label, show_label=True), gr.Number(label="Predicted Count", visible=True, value=boxes.shape[0]))
 
 def remove_label(image):