logging

handler.py

@@ -1,5 +1,6 @@
 import torch
 from transformers import pipeline
+from transformers.utils import logging
 from PIL import Image
 import requests
 import numpy as np
@@ -7,16 +8,19 @@ from cv2 import imencode
 from base64 import b64encode
 import time
 
+logging.set_verbosity_info()
+
 class EndpointHandler():
   def __init__(self, path=""):
     self.device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f"Using device: {self.device}")
+    self.logger = logging.get_logger("transformers")
+    self.logger.info(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)
-    print(f"got request for {inputs}")
+    self.logger.info(f"got request for {inputs}")
     raw_image = Image.open(requests.get(inputs, stream=True).raw).convert("RGB")
     
     with torch.no_grad():