For users to load in one key.

BiRefNet_github/models/birefnet.py

@@ -114,7 +114,7 @@ class BiRefNet(
     def forward(self, x):
         scaled_preds, class_preds = self.forward_ori(x)
         class_preds_lst = [class_preds]
-        return [scaled_preds, class_preds_lst] if self.training else scaled_preds
+        return [scaled_preds, class_preds_lst] if self.training and 0 else scaled_preds
 
 
 class Decoder(nn.Module):

config.json

@@ -4,7 +4,7 @@
     "BiRefNet"
   ],
   "auto_map": {
-    "AutoModelForImageSegmentation": "BiRefNet.models.birefnet.BiRefNet"
+    "AutoModelForImageSegmentation": "BiRefNet_github.models.birefnet.BiRefNet"
   },
   "custom_pipelines": {
     "image-segmentation": {