bytes

handler.py

@@ -5,13 +5,16 @@ import requests
 import numpy as np
 from cv2 import imencode
 from base64 import b64encode
+import time
 
 class EndpointHandler():
   def __init__(self, path=""):
     self.device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {self.device}")
     self.generator = pipeline("mask-generation", model="facebook/sam-vit-large", device=self.device)
 
   def __call__(self, data):
+    start = time.time()
     inputs = data.pop("inputs", data)
     raw_image = Image.open(requests.get(inputs, stream=True).raw).convert("RGB")
     
@@ -31,5 +34,7 @@ class EndpointHandler():
           "mask": encoded_mask,
           "label": ""
       })
+
+    end = time.time()
     
-    return data
+    return { "data": data, "time": end - start }