Align lib_name as birefnet and add inference endpoint option.

README.md

@@ -1,5 +1,5 @@
 ---
-library_name: BiRefNet-legacy
+library_name: birefnet
 tags:
 - background-removal
 - mask-generation

birefnet.py

@@ -1995,7 +1995,8 @@ class BiRefNet(
 ):
     config_class = BiRefNetConfig
     def __init__(self, bb_pretrained=True, config=BiRefNetConfig()):
-        super(BiRefNet, self).__init__()
+        super(BiRefNet, self).__init__(config)
+        bb_pretrained = config.bb_pretrained
         self.config = Config()
         self.epoch = 1
         self.bb = build_backbone(self.config.bb, pretrained=bb_pretrained)

handler.py

@@ -62,6 +62,7 @@ class ImagePreprocessor():
 
 usage_to_weights_file = {
     'General': 'BiRefNet',
+    'General-HR': 'BiRefNet_HR',
     'General-Lite': 'BiRefNet_lite',
     'General-Lite-2K': 'BiRefNet_lite-2K',
     'General-reso_512': 'BiRefNet-reso_512',
@@ -82,9 +83,12 @@ if usage in ['General-Lite-2K']:
     resolution = (2560, 1440)
 elif usage in ['General-reso_512']:
     resolution = (512, 512)
+elif usage in ['General-HR']:
+    resolution = (2048, 2048)
 else:
     resolution = (1024, 1024) 
 
+half_precision = True
 
 class EndpointHandler():
     def __init__(self, path=''):
@@ -93,6 +97,8 @@ class EndpointHandler():
         )
         self.birefnet.to(device)
         self.birefnet.eval()
+        if half_precision:
+            self.birefnet.half()
 
     def __call__(self, data: Dict[str, Any]):
         """
@@ -122,7 +128,7 @@ class EndpointHandler():
 
         # Prediction
         with torch.no_grad():
-            preds = self.birefnet(image_proc.to(device))[-1].sigmoid().cpu()
+            preds = self.birefnet(image_proc.to(device).half() if half_precision else image_proc.to(device))[-1].sigmoid().cpu()
         pred = preds[0].squeeze()
 
         # Show Results