From 7fa2d0d357735c117bbe9ed07d4432567eb67d9f Mon Sep 17 00:00:00 2001 From: bigcat88 Date: Fri, 20 Dec 2024 14:54:47 +0200 Subject: [PATCH] use "comfy.model_management.get_torch_device()" instead if "torch.cuda.is_available" --- py/imagefunc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/py/imagefunc.py b/py/imagefunc.py index cb276d2..9c74e2a 100644 --- a/py/imagefunc.py +++ b/py/imagefunc.py @@ -1453,7 +1453,7 @@ def create_mask_from_color_tensor(image:Image, color:str, tolerance:int=0) -> Im def load_RMBG_model(): from .briarmbg import BriaRMBG current_directory = os.path.dirname(os.path.abspath(__file__)) - device = "cuda" if torch.cuda.is_available() else "cpu" + device = comfy.model_management.get_torch_device() net = BriaRMBG() model_path = "" try: