Add inference endpoint feature.

handler.py

@@ -75,9 +75,7 @@ usage_to_weights_file = {
 }
 
 usage = 'General'
-birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file[usage])), trust_remote_code=True)
-birefnet.to(device)
-birefnet.eval()
+model_repo = '/'.join(('zhengpeng7', usage_to_weights_file[usage]))
 
 # Set resolution
 if usage in ['General-Lite-2K']:
@@ -91,9 +89,10 @@ else:
 class EndpointHandler():
     def __init__(self, path=""):
         self.birefnet = AutoModelForImageSegmentation.from_pretrained(
-            "ZhengPeng7/BiRefNet", trust_remote_code=True
+            model_repo, trust_remote_code=True
         )
         self.birefnet.to(device)
+        self.birefnet.eval()
 
     def __call__(self, data: Dict[str, Any]):
         """
@@ -123,7 +122,7 @@ class EndpointHandler():
 
         # Prediction
         with torch.no_grad():
-            preds = birefnet(image_proc.to(device))[-1].sigmoid().cpu()
+            preds = self.birefnet(image_proc.to(device))[-1].sigmoid().cpu()
         pred = preds[0].squeeze()
 
         # Show Results