fix auto device

This commit is contained in:
yang shawn
2024-03-22 10:33:54 +08:00
parent 0c3066f4e9
commit 1aca549818
+9 -8
View File
@@ -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)