Enabled ZeroGPU

app.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 import copy
 import random
@@ -92,6 +93,7 @@ def get_args_parser():
     parser.add_argument("--amp", action="store_true", help="Train with mixed precision")
     return parser
 
+@spaces.GPU
 # Get counting model.
 def build_model_and_transforms(args):
     normalize = T.Compose(
@@ -194,6 +196,7 @@ def get_ind_to_filter(text, word_ids, keywords):
 
     return inds_to_filter
 
+@spaces.GPU
 def count(image, text, prompts, state):
     print("state: " + str(state))
     keywords = "" # do not handle this for now
@@ -282,6 +285,7 @@ def count(image, text, prompts, state):
         out_label = "Nothing specified to detect."
     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):
     keywords = "" # do not handle this for now
     # Handle no prompt case.