From 172e873e926a909d7dc65acd0910071487119ca9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E9=9B=AA=E5=B3=B0?= <lldacing@163.com>
Date: Sun, 10 Nov 2024 00:38:43 +0800
Subject: [PATCH] fix cpu error

---
 birefnetNode.py | 3 ++-
 pyproject.toml  | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/birefnetNode.py b/birefnetNode.py
index 69ceaa9..fb88da5 100644
--- 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
index 054f78d..d1ab829 100644
--- 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"]