Add inference endpoint feature.

README.md

@@ -7,7 +7,7 @@ tags:
 - pytorch_model_hub_mixin
 - model_hub_mixin
 repo_url: https://github.com/ZhengPeng7/BiRefNet-legacy
-pipeline_tag: image-to-image
+pipeline_tag: image-segmentation
 ---
 <h1 align="center">Bilateral Reference for High-Resolution Dichotomous Image Segmentation</h1>
 

handler.py

@@ -74,8 +74,8 @@ usage_to_weights_file = {
     'General-legacy': 'BiRefNet-legacy'
 }
 
+# Choose the version of BiRefNet here.
 usage = 'General'
-model_repo = '/'.join(('zhengpeng7', usage_to_weights_file[usage]))
 
 # Set resolution
 if usage in ['General-Lite-2K']:
@@ -87,9 +87,9 @@ else:
 
 
 class EndpointHandler():
-    def __init__(self, path=""):
+    def __init__(self):
         self.birefnet = AutoModelForImageSegmentation.from_pretrained(
-            model_repo, trust_remote_code=True
+            '/'.join(('zhengpeng7', usage_to_weights_file[usage])), trust_remote_code=True
         )
         self.birefnet.to(device)
         self.birefnet.eval()