Skip to content

Commit

Permalink
fix:fail to load weight when do not configure birefnet path
Browse files Browse the repository at this point in the history
  • Loading branch information
刘雪峰 committed Sep 13, 2024
1 parent 930be71 commit 858acf9
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 8 deletions.
12 changes: 12 additions & 0 deletions __init__.py
Original file line number Diff line number Diff line change
@@ -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}
Expand Down
8 changes: 0 additions & 8 deletions birefnetNode.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]

Expand Down

0 comments on commit 858acf9

Please sign in to comment.