fix: error AutoDownloadBiRefNetModel node's name

lldacing · Sep 6, 2024 · a6617a0 · a6617a0
1 parent f47fbcb
commit a6617a0
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 8 deletions.
diff --git a/birefnet/config.py b/birefnet/config.py
@@ -1,9 +1,10 @@
 import os
 import math
 
-CUR_DIR = os.path.dirname(__file__)
+# CUR_DIR = os.path.dirname(__file__)
 
-class Config():
+
+class Config:
     def __init__(self, bb_index: int = 6) -> None:
         # PATH settings
         # Make up your file system as: SYS_HOME_DIR/codes/dis/BiRefNet, SYS_HOME_DIR/datasets/dis/xx, SYS_HOME_DIR/weights/xx

diff --git a/birefnetNode.py b/birefnetNode.py
@@ -6,13 +6,10 @@
 import comfy
 from comfy import model_management
 import folder_paths
-from birefnet.config import Config
 from birefnet.models.birefnet import BiRefNet
 from birefnet.utils import check_state_dict
 from .util import tensor_to_pil, apply_mask_to_image
 
-config = Config()
-
 deviceType = model_management.get_torch_device().type
 
 models_dir_key = "birefnet"
@@ -193,13 +190,13 @@ def rem_bg(self, model, images):
 
 
 NODE_CLASS_MAPPINGS = {
-    "LoadBiRefNetModelByName": AutoDownloadBiRefNetModel,
+    "AutoDownloadBiRefNetModel": AutoDownloadBiRefNetModel,
     "LoadRembgByBiRefNetModel": LoadRembgByBiRefNetModel,
     "RembgByBiRefNet": RembgByBiRefNet,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "LoadBiRefNetModelByName": "LoadBiRefNetModelByName",
+    "AutoDownloadBiRefNetModel": "AutoDownloadBiRefNetModel",
     "LoadRembgByBiRefNetModel": "LoadRembgByBiRefNetModel",
     "RembgByBiRefNet": "RembgByBiRefNet",
 }
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.0"
+version = "1.0.1"
 license = {file = "LICENSE"}
 dependencies = ["numpy<2", "opencv-python", "scipy", "timm"]