fix cpu error

lldacing · Nov 9, 2024 · 172e873 · 172e873
1 parent 47fd1f8
commit 172e873
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/birefnetNode.py b/birefnetNode.py
@@ -178,6 +178,7 @@ def INPUT_TYPES(cls):
 
     def rem_bg(self, model, images):
         model, version = model
+        model_device_type = next(model.parameters()).device.type
         _images = []
         _masks = []
 
@@ -191,7 +192,7 @@ def rem_bg(self, model, images):
                 im_tensor = proc_img(pil_image).unsqueeze(0)
 
             with torch.no_grad():
-                mask = model(im_tensor.to(deviceType))[-1].sigmoid().cpu()
+                mask = model(im_tensor.to(model_device_type))[-1].sigmoid().cpu()
 
             # 遮罩大小需还原为与原图一致
             mask = comfy.utils.common_upscale(mask, w, h, 'bilinear', "disabled")

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui_birefnet_ll"
 description = "Sync with version of BiRefNet. NODES:AutoDownloadBiRefNetModel, LoadRembgByBiRefNetModel, RembgByBiRefNet."
-version = "1.0.4"
+version = "1.0.5"
 license = {file = "LICENSE"}
 dependencies = ["numpy<2", "opencv-python", "scipy", "timm"]