fix auto device
This commit is contained in:
+9
-8
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user