From 1aca549818596ea31f7e50c26ff77c88e9439f13 Mon Sep 17 00:00:00 2001
From: yang shawn <shawnyangxu@gmail.com>
Date: Fri, 22 Mar 2024 10:33:54 +0800
Subject: [PATCH] fix auto device

---
 BiRefNet_node.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/BiRefNet_node.py b/BiRefNet_node.py
index 2227938..b65f501 100644
--- a/BiRefNet_node.py
+++ b/BiRefNet_node.py
@@ -44,14 +44,6 @@ 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')
@@ -97,6 +89,15 @@ class BiRefNet_node:
     CATEGORY = "Fooocus"
 
     def matting(self, image, device):
+        # process auto deivce
+        if device == "auto":
+            if torch.backends.mps.is_available():
+                device = "mps"
+            elif torch.cuda.is_available():
+                device = "cuda"
+            else:
+                device = "cpu"
+
         if not self.ready:
             weight_path = os.path.join(models_dir, "BiRefNet", "BiRefNet-ep480.pth")
             self.load(weight_path, device=device)