trying out ZeroGPU

model_yolov5.py

@@ -1,16 +1,18 @@
+import spaces
 import numpy as np
 import yolov5
 from yolov5.utils.plots import Annotator, colors
 from huggingface_hub import get_token
 
 
+@spaces.GPU
 def load_model(model_path, img_size=640):
     """Load model from HuggingFace Hub."""
     model = yolov5.load(model_path, hf_token=get_token())
     model.img_size = img_size  # add img_size attribute
     return model
 
-
+@spaces.GPU
 def inference(model, image):
     """Run inference on image and return annotated image."""
     results = model(image, size=model.img_size)

requirements.txt

@@ -1,2 +1,3 @@
+spaces
 yolov5
 pillow