Merge pull request #460 from Visionatrix/main

rmbg1.4: use "comfy.get_torch_device()" instead if "torch.cuda.is_available"
2024-12-20 22:38:21 +08:00
parent c6ac429ef6 7fa2d0d357
commit 666541eb06
1 changed files with 1 additions and 1 deletions
@@ -1478,7 +1478,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: