From a6617a0e27348e3929440293c80e4a626509dad0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E9=9B=AA=E5=B3=B0?= <lldancing@163.com>
Date: Fri, 6 Sep 2024 14:14:11 +0800
Subject: [PATCH] fix: error AutoDownloadBiRefNetModel node's name

---
 birefnet/config.py | 5 +++--
 birefnetNode.py    | 7 ++-----
 pyproject.toml     | 2 +-
 3 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/birefnet/config.py b/birefnet/config.py
index 4ed688d..598c245 100644
--- 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
index 5ac6240..6a6a7d4 100644
--- 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
index f0277f0..f4f786b 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.0"
+version = "1.0.1"
 license = {file = "LICENSE"}
 dependencies = ["numpy<2", "opencv-python", "scipy", "timm"]