fix auto device

2024-03-22 10:23:17 +08:00
parent 50e248b929
commit 0c3066f4e9
1 changed files with 8 additions and 0 deletions
@@ -44,6 +44,14 @@ class BiRefNet_node:
        self.ready = False

    def load(self, weight_path, device, verbose=False):
+        # process auto
+        if device == "auto":
+            if torch.backends.mps.is_available():
+                device = "mps"
+            elif torch.cuda.is_available():
+                device = "cuda"
+            else:
+                device = "cpu"
        # load model
        self.model = BiRefNet()
        state_dict = torch.load(weight_path, map_location='cpu')