feat: initial project setup
This commit is contained in:
@@ -0,0 +1,262 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""模型代码定位与 ComfyUI 解耦层。
|
||||
|
||||
本项目把原来 ComfyUI 节点 ``comfyui_birefnet_ll`` 的模型代码**内联**到了
|
||||
``vendor/comfyui_birefnet_ll/``，因此默认情况下完全自包含，不再依赖外部 ComfyUI 安装。
|
||||
|
||||
定位优先级（``find_node_dir``）：
|
||||
1. 命令行 ``--node-dir``
|
||||
2. 环境变量 ``BIREFNET_NODE_DIR``
|
||||
3. **项目内 ``vendor/comfyui_birefnet_ll``（默认，开箱即用）**
|
||||
4. ``config.json`` 中的 ``node_dir`` 字段（兼容旧配置）
|
||||
5. 自动扫描本机 ComfyUI 安装（便于跟随上游节点升级）
|
||||
|
||||
模型包在 import 阶段依赖 ComfyUI 的 ``folder_paths``，但只用于「按文件名查权重路径」
|
||||
（即骨干网络预训练权重）。本项目始终以 ``bb_pretrained=False`` 构建骨干网络，不会真正
|
||||
读取这类权重，因此这里注入一个最小实现的 ``folder_paths`` 垫片即可脱离 ComfyUI 运行。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
from typing import Dict, Iterable, List, Optional, Tuple
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 路径常量
|
||||
# --------------------------------------------------------------------------- #
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||||
|
||||
#: 项目内自带的模型代码（默认来源，随项目一起分发）
|
||||
VENDOR_DIR = PROJECT_ROOT / "vendor" / "comfyui_birefnet_ll"
|
||||
|
||||
#: 项目内自带的权重目录（默认来源）
|
||||
LOCAL_MODEL_DIR = PROJECT_ROOT / "models"
|
||||
|
||||
#: 环境变量：手动指定模型代码目录
|
||||
NODE_DIR_ENV = "BIREFNET_NODE_DIR"
|
||||
|
||||
#: 节点仓库名，用于兜底扫描外部 ComfyUI 安装
|
||||
NODE_REPO_NAME = "comfyui_birefnet_ll"
|
||||
|
||||
#: 旧版架构权重文件名（与 birefnetNode.py 保持一致）
|
||||
OLD_MODEL_FILES: Tuple[str, ...] = ("BiRefNet-DIS_ep580.pth", "BiRefNet-ep480.pth")
|
||||
|
||||
#: 判定「是有效模型代码目录」的标志文件
|
||||
_MARKER = ("birefnet", "models", "birefnet.py")
|
||||
|
||||
_BUNDLED_PYTHON = PROJECT_ROOT / "python" / "python.exe"
|
||||
|
||||
_INSTALL_HINT = (
|
||||
"缺少运行依赖。本项目自带独立 Python 运行时，请使用：\n"
|
||||
f' "{_BUNDLED_PYTHON}" Webui.py\n'
|
||||
"或双击 run.bat。\n"
|
||||
"若想用自己的解释器，请先安装依赖：pip install -r requirements.txt"
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 模型代码目录定位
|
||||
# --------------------------------------------------------------------------- #
|
||||
def _is_code_dir(path: Path) -> bool:
|
||||
"""目录内是否含 ``birefnet/models/birefnet.py``。"""
|
||||
try:
|
||||
return path.joinpath(*_MARKER).is_file()
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
def _scan_external_node_dirs() -> List[Path]:
|
||||
"""扫描本机常见的 ComfyUI 安装位置（仅在需要跟随上游升级时才会命中）。"""
|
||||
found: List[Path] = []
|
||||
home = Path.home()
|
||||
roots: List[Path] = [home]
|
||||
for letter in "CDEFGHIJ":
|
||||
roots.append(Path(f"{letter}:/"))
|
||||
patterns = (
|
||||
"ComfyUI/custom_nodes/" + NODE_REPO_NAME,
|
||||
"ComfyUI*/ComfyUI/custom_nodes/" + NODE_REPO_NAME,
|
||||
"ComfyUI*/custom_nodes/" + NODE_REPO_NAME,
|
||||
"ComfyUI*/ComfyUI/custom_nodes/ComfyUI_BiRefNet_ll",
|
||||
"ComfyUI*/custom_nodes/ComfyUI_BiRefNet_ll",
|
||||
"*/*/ComfyUI/custom_nodes/" + NODE_REPO_NAME,
|
||||
)
|
||||
for root in roots:
|
||||
if not root.exists():
|
||||
continue
|
||||
for pat in patterns:
|
||||
try:
|
||||
found.extend(sorted(root.glob(pat)))
|
||||
except (OSError, ValueError):
|
||||
continue
|
||||
return found
|
||||
|
||||
|
||||
def _dedup(paths: Iterable[Path]) -> List[Path]:
|
||||
seen, out = set(), []
|
||||
for p in paths:
|
||||
try:
|
||||
key = str(p.resolve()).lower() if p.exists() else str(p).lower()
|
||||
except OSError:
|
||||
key = str(p).lower()
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
out.append(p)
|
||||
return out
|
||||
|
||||
|
||||
def candidate_code_dirs(config_hint: Optional[str] = None) -> List[Path]:
|
||||
"""按优先级生成候选模型代码目录。"""
|
||||
cands: List[Path] = []
|
||||
|
||||
env = os.environ.get(NODE_DIR_ENV)
|
||||
if env:
|
||||
cands.append(Path(env))
|
||||
|
||||
# 项目自带（默认命中）
|
||||
cands.append(VENDOR_DIR)
|
||||
|
||||
# 旧配置中记录的路径（兼容早期版本写死的 ComfyUI 节点路径）
|
||||
if config_hint:
|
||||
cands.append(Path(config_hint))
|
||||
|
||||
# 外部 ComfyUI 安装（兜底）
|
||||
cands.append(PROJECT_ROOT.parent / NODE_REPO_NAME)
|
||||
cands.extend(_scan_external_node_dirs())
|
||||
|
||||
return _dedup(cands)
|
||||
|
||||
|
||||
def find_node_dir(explicit: Optional[str] = None, config_hint: Optional[str] = None) -> Path:
|
||||
"""定位模型代码目录（需含 ``birefnet/models/birefnet.py``）。
|
||||
|
||||
正常情况下会直接命中项目内的 ``vendor/comfyui_birefnet_ll``。
|
||||
"""
|
||||
cands: List[Path] = []
|
||||
if explicit:
|
||||
cands.append(Path(explicit))
|
||||
cands.extend(candidate_code_dirs(config_hint))
|
||||
|
||||
for cand in cands:
|
||||
if _is_code_dir(cand):
|
||||
try:
|
||||
return cand.resolve()
|
||||
except OSError:
|
||||
return cand
|
||||
|
||||
raise FileNotFoundError(
|
||||
"未找到 BiRefNet 模型代码目录（应包含 birefnet/models/birefnet.py）。\n"
|
||||
f"默认位置：{VENDOR_DIR}\n"
|
||||
"若该目录缺失或被误删，可从 ComfyUI 节点 comfyui_birefnet_ll 重新拷贝，或用以下方式指定：\n"
|
||||
" 1) 启动参数 --node-dir <路径>\n"
|
||||
f" 2) 环境变量 {NODE_DIR_ENV}=<路径>\n"
|
||||
" 3) config.json 中的 node_dir 字段\n"
|
||||
"已尝试的候选目录：\n - " + "\n - ".join(str(c) for c in cands[:8])
|
||||
)
|
||||
|
||||
|
||||
def default_model_dirs(code_dir: Optional[Path] = None) -> List[Path]:
|
||||
"""返回默认权重搜索目录：项目内 ``models/`` 优先，其次外部目录。"""
|
||||
dirs: List[Path] = [LOCAL_MODEL_DIR]
|
||||
|
||||
# 若用户显式指向了外部 ComfyUI 节点，则顺带搜索对应的 models/BiRefNet
|
||||
if code_dir is not None:
|
||||
p = Path(code_dir)
|
||||
try:
|
||||
if p.resolve() != VENDOR_DIR.resolve():
|
||||
dirs.append(p.parent.parent / "models" / "BiRefNet")
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
return _dedup(dirs)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# folder_paths 垫片
|
||||
# --------------------------------------------------------------------------- #
|
||||
class _FolderPathsShim(types.ModuleType):
|
||||
"""仅实现 birefnet/config.py 用到的接口，一律返回 None（不查权重）。"""
|
||||
|
||||
def __init__(self, model_dirs: Iterable[Path]) -> None:
|
||||
super().__init__("folder_paths")
|
||||
self.models_dir = str(next(iter(model_dirs), Path("models")))
|
||||
self.folder_names_and_paths: Dict[str, Tuple[List[str], set]] = {
|
||||
"birefnet": ([self.models_dir], {".pt", ".pth", ".safetensors", ".ckpt"})
|
||||
}
|
||||
self.supported_pt_extensions = {".pt", ".pth", ".safetensors", ".ckpt"}
|
||||
|
||||
def get_folder_paths(self, folder_name: str) -> List[str]:
|
||||
return list(self.folder_names_and_paths.get(folder_name, ([], None))[0])
|
||||
|
||||
def get_full_path(self, folder_name: str, filename: str) -> Optional[str]:
|
||||
return None
|
||||
|
||||
def get_filename_list(self, folder_name: str) -> List[str]:
|
||||
return []
|
||||
|
||||
def add_model_folder_path(self, folder_name: str, full_folder_path: str) -> None:
|
||||
if folder_name not in self.folder_names_and_paths:
|
||||
self.folder_names_and_paths[folder_name] = ([], set())
|
||||
self.folder_names_and_paths[folder_name][0].append(str(full_folder_path))
|
||||
|
||||
def filter_files_extensions(self, files, extensions): # pragma: no cover
|
||||
return [f for f in files if os.path.splitext(f)[1] in extensions]
|
||||
|
||||
|
||||
def ensure_folder_paths(model_dirs: Iterable[Path]) -> bool:
|
||||
"""确保 ``folder_paths`` 可用。返回 True 表示用了真实实现（即运行在 ComfyUI 内）。"""
|
||||
if "folder_paths" in sys.modules:
|
||||
return not isinstance(sys.modules["folder_paths"], _FolderPathsShim)
|
||||
try: # 若真的在 ComfyUI 环境内运行，优先用真实实现
|
||||
importlib.import_module("folder_paths")
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
sys.modules["folder_paths"] = _FolderPathsShim(model_dirs)
|
||||
return False
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 模型类导入
|
||||
# --------------------------------------------------------------------------- #
|
||||
_cached: Dict[str, object] = {}
|
||||
|
||||
|
||||
def load_model_classes(node_dir: Path, model_dirs: Iterable[Path]) -> Dict[str, object]:
|
||||
"""导入并返回 ``{"BiRefNet":..., "OldBiRefNet":..., "check_state_dict":..., "shim": bool}``。
|
||||
|
||||
结果会被缓存，重复调用不会重复导入。
|
||||
"""
|
||||
if _cached:
|
||||
return _cached
|
||||
|
||||
node_dir = Path(node_dir).resolve()
|
||||
shim = ensure_folder_paths(model_dirs)
|
||||
|
||||
p = str(node_dir)
|
||||
if p not in sys.path:
|
||||
sys.path.insert(0, p)
|
||||
|
||||
try:
|
||||
from birefnet.models.birefnet import BiRefNet
|
||||
from birefnet.utils import check_state_dict
|
||||
except ImportError as exc: # pragma: no cover - 环境问题
|
||||
raise ImportError(f"导入 birefnet 模型包失败（{node_dir}）：{exc}\n{_INSTALL_HINT}") from exc
|
||||
|
||||
try:
|
||||
from birefnet_old.models.birefnet import BiRefNet as OldBiRefNet
|
||||
except Exception: # 旧包缺失不影响新模型使用
|
||||
OldBiRefNet = None # type: ignore[assignment]
|
||||
|
||||
_cached.update(
|
||||
BiRefNet=BiRefNet,
|
||||
OldBiRefNet=OldBiRefNet,
|
||||
check_state_dict=check_state_dict,
|
||||
node_dir=node_dir,
|
||||
used_folder_paths_shim=shim,
|
||||
)
|
||||
return _cached
|
||||
Reference in New Issue
Block a user