From 858acf929381c667e502f068b33507fdac1974f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=88=98=E9=9B=AA=E5=B3=B0?= Date: Fri, 13 Sep 2024 15:32:33 +0800 Subject: [PATCH] fix:fail to load weight when do not configure birefnet path --- __init__.py | 12 ++++++++++++ birefnetNode.py | 8 -------- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/__init__.py b/__init__.py index 5abad98..2adb1ab 100644 --- a/__init__.py +++ b/__init__.py @@ -1,12 +1,24 @@ import os import sys +import folder_paths + # 获取当前目录的父目录的父目录 parent_dir = os.path.dirname(os.path.abspath(__file__)) # 添加父目录的父目录到系统路径 sys.path.insert(0, parent_dir) +models_dir_key = "birefnet" +models_dir_default = os.path.join(folder_paths.models_dir, "BiRefNet") +if models_dir_key not in folder_paths.folder_names_and_paths: + folder_paths.folder_names_and_paths[models_dir_key] = ( + [os.path.join(folder_paths.models_dir, "BiRefNet")], folder_paths.supported_pt_extensions) +else: + if not os.path.exists(models_dir_default): + os.makedirs(models_dir_default, exist_ok=True) + folder_paths.add_model_folder_path(models_dir_key, models_dir_default) + from . import birefnetNode NODE_CLASS_MAPPINGS = {**birefnetNode.NODE_CLASS_MAPPINGS} diff --git a/birefnetNode.py b/birefnetNode.py index 7a33d54..cd5bfaf 100644 --- a/birefnetNode.py +++ b/birefnetNode.py @@ -14,14 +14,6 @@ deviceType = model_management.get_torch_device().type models_dir_key = "birefnet" -models_dir_default = os.path.join(folder_paths.models_dir, "BiRefNet") -if models_dir_key not in folder_paths.folder_names_and_paths: - folder_paths.folder_names_and_paths[models_dir_key] = ( - [os.path.join(folder_paths.models_dir, "BiRefNet")], folder_paths.supported_pt_extensions) -else: - if not os.path.exists(models_dir_default): - os.makedirs(models_dir_default, exist_ok=True) - folder_paths.add_model_folder_path(models_dir_key, models_dir_default) models_path_default = folder_paths.get_folder_paths(models_dir_key)[0]