358 lines
14 KiB
Python
358 lines
14 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""BiRefNet WebUI 启动入口。
|
||
|
||
本项目**自包含**：模型代码、权重、Python 运行时都在项目目录内，不再需要 ComfyUI。
|
||
|
||
F:\\BiRefNet_WebUI\\
|
||
├─ Webui.py
|
||
├─ python\\ 自带 Python 3.12 运行时（torch + CUDA）
|
||
├─ models\\ BiRefNet 权重（*.safetensors）
|
||
├─ vendor\\comfyui_birefnet_ll\\ 模型代码
|
||
└─ birefnet_web\\ web\\ outputs\\
|
||
|
||
启动方式（任选其一）：
|
||
|
||
run.bat 双击即可
|
||
python\\python.exe Webui.py 用自带运行时
|
||
python Webui.py 用系统 Python（缺依赖时会自动切到自带运行时）
|
||
|
||
常用参数：
|
||
--port 7861 监听端口
|
||
--host 0.0.0.0 允许局域网访问
|
||
--device cpu 强制 CPU 推理
|
||
--dtype float32 强制单精度（默认 GPU 上 fp32 权重 + fp16 autocast）
|
||
--model-dir D:\\models 追加模型目录（可重复）
|
||
--node-dir <path> 指定模型代码目录（默认用项目内 vendor/）
|
||
--no-browser 启动后不自动打开浏览器
|
||
--check 仅自检环境后退出
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import sys
|
||
import threading
|
||
import time
|
||
import webbrowser
|
||
from pathlib import Path
|
||
from typing import Dict, List, Optional
|
||
|
||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||
CONFIG_PATH = PROJECT_ROOT / "config.json"
|
||
|
||
# 内置 Python 带有 python3xx._pth（隔离模式），脚本所在目录不会自动进入 sys.path，
|
||
# 这里显式补上，保证无论用哪个解释器都能 import 到本项目包。
|
||
if str(PROJECT_ROOT) not in sys.path:
|
||
sys.path.insert(0, str(PROJECT_ROOT))
|
||
|
||
#: 项目自带的独立 Python 运行时
|
||
BUNDLED_PYTHON = PROJECT_ROOT / "python" / "python.exe"
|
||
|
||
#: 防止自动切换解释器时无限递归
|
||
_RELAUNCH_FLAG = "BIREFNET_RELAUNCHED"
|
||
|
||
#: 运行必需的三方包 -> pip 包名
|
||
_REQUIRED = (
|
||
("torch", "torch"),
|
||
("torchvision", "torchvision"),
|
||
("numpy", "numpy"),
|
||
("PIL", "pillow"),
|
||
("safetensors", "safetensors"),
|
||
("timm", "timm"),
|
||
("einops", "einops"),
|
||
("kornia", "kornia"),
|
||
)
|
||
|
||
CONFIG_DEFAULTS: Dict[str, object] = {
|
||
"host": "127.0.0.1",
|
||
"port": 7861,
|
||
"node_dir": "",
|
||
"model_dirs": [],
|
||
"output_dir": "outputs",
|
||
"device": "auto",
|
||
"dtype": "auto",
|
||
"max_cached_models": 1,
|
||
"max_tasks": 50,
|
||
"open_browser": True,
|
||
}
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 依赖自检 / 解释器自动切换
|
||
# --------------------------------------------------------------------------- #
|
||
def missing_packages() -> List[str]:
|
||
"""返回缺失的 pip 包名列表。"""
|
||
missing: List[str] = []
|
||
for mod, pkg in _REQUIRED:
|
||
try:
|
||
__import__(mod)
|
||
except Exception:
|
||
missing.append(pkg)
|
||
return missing
|
||
|
||
|
||
def bundled_python() -> Optional[Path]:
|
||
"""项目自带运行时是否可用。"""
|
||
return BUNDLED_PYTHON if BUNDLED_PYTHON.is_file() else None
|
||
|
||
|
||
def relaunch_with_bundled_python(missing: List[str]) -> None:
|
||
"""当前解释器缺依赖、而项目自带运行时可用时，直接把进程换成自带运行时。"""
|
||
if not missing or os.environ.get(_RELAUNCH_FLAG):
|
||
return
|
||
exe = bundled_python()
|
||
if exe is None:
|
||
return
|
||
try:
|
||
if Path(sys.executable).resolve() == exe.resolve():
|
||
return
|
||
except OSError:
|
||
pass
|
||
os.environ[_RELAUNCH_FLAG] = "1"
|
||
print(f"[环境] 当前解释器缺少 {', '.join(missing)}，自动切换到项目自带运行时：{exe}", flush=True)
|
||
try:
|
||
os.execv(str(exe), [str(exe), str(Path(__file__).resolve()), *sys.argv[1:]])
|
||
except OSError as exc: # 切换失败就继续走下面的报错分支
|
||
print(f"[环境] 切换失败：{exc}", file=sys.stderr, flush=True)
|
||
|
||
|
||
def preflight() -> None:
|
||
"""尽早给出人话的依赖缺失提示，而不是一堆 ImportError 堆栈。"""
|
||
missing = missing_packages()
|
||
if missing:
|
||
relaunch_with_bundled_python(missing)
|
||
if missing:
|
||
print("缺少依赖：" + ", ".join(missing), file=sys.stderr)
|
||
if bundled_python() is None:
|
||
print(
|
||
f"项目自带运行时不存在（{BUNDLED_PYTHON}）。\n"
|
||
"请从 ComfyUI 便携版的 python_embeded 目录复制一份到项目的 python\\ 下，"
|
||
"或用你自己的解释器安装依赖：",
|
||
file=sys.stderr,
|
||
)
|
||
print(f' "{sys.executable}" -m pip install -r requirements.txt', file=sys.stderr)
|
||
sys.exit(2)
|
||
if not _has_cv2():
|
||
print("[提示] 未检测到 opencv-python，前景精修将退化为高斯模糊实现（结果略有差异）")
|
||
|
||
|
||
def _has_cv2() -> bool:
|
||
try:
|
||
import cv2 # noqa: F401
|
||
|
||
return True
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 配置
|
||
# --------------------------------------------------------------------------- #
|
||
def load_config(path: Path) -> Dict[str, object]:
|
||
config = dict(CONFIG_DEFAULTS)
|
||
if path.is_file():
|
||
try:
|
||
data = json.loads(path.read_text(encoding="utf-8"))
|
||
if isinstance(data, dict):
|
||
config.update(data)
|
||
except (json.JSONDecodeError, OSError) as exc:
|
||
print(f"[警告] 读取 {path} 失败，使用默认配置：{exc}")
|
||
return config
|
||
|
||
|
||
def save_config(path: Path, config: Dict[str, object]) -> None:
|
||
try:
|
||
path.write_text(json.dumps(config, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
except OSError as exc: # pragma: no cover
|
||
print(f"[警告] 写入 {path} 失败：{exc}")
|
||
|
||
|
||
def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(
|
||
prog="Webui.py",
|
||
description="BiRefNet WebUI —— 自包含的网页版智能抠图服务",
|
||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||
)
|
||
parser.add_argument("--host", default=None, help="监听地址，0.0.0.0 可局域网访问")
|
||
parser.add_argument("--port", type=int, default=None, help="监听端口")
|
||
parser.add_argument("--device", default=None, choices=["auto", "cpu", "cuda"], help="推理设备")
|
||
parser.add_argument("--dtype", default=None, choices=["auto", "float32", "float16", "bfloat16"], help="推理精度")
|
||
parser.add_argument("--model-dir", action="append", default=None, help="追加模型目录，可重复指定")
|
||
parser.add_argument("--node-dir", default=None, help="模型代码目录（默认用项目内 vendor/）")
|
||
parser.add_argument("--output-dir", default=None, help="结果输出目录")
|
||
parser.add_argument("--max-cached-models", type=int, default=None, help="显存中最多缓存的模型数量")
|
||
parser.add_argument("--max-tasks", type=int, default=None, help="内存中保留的任务记录数量")
|
||
parser.add_argument("--config", default=str(CONFIG_PATH), help="配置文件路径")
|
||
parser.add_argument("--no-browser", action="store_true", help="启动后不自动打开浏览器")
|
||
parser.add_argument("--save-config", action="store_true", help="把本次参数写回配置文件")
|
||
parser.add_argument("--check", action="store_true", help="仅做环境自检并打印可用模型，然后退出")
|
||
return parser.parse_args(argv)
|
||
|
||
|
||
def merge_config(args: argparse.Namespace) -> Dict[str, object]:
|
||
config_path = Path(args.config).expanduser().resolve()
|
||
config = load_config(config_path)
|
||
|
||
if args.host is not None:
|
||
config["host"] = args.host
|
||
if args.port is not None:
|
||
config["port"] = args.port
|
||
if args.device is not None:
|
||
config["device"] = args.device
|
||
if args.dtype is not None:
|
||
config["dtype"] = args.dtype
|
||
if args.output_dir is not None:
|
||
config["output_dir"] = args.output_dir
|
||
if args.max_cached_models is not None:
|
||
config["max_cached_models"] = args.max_cached_models
|
||
if args.max_tasks is not None:
|
||
config["max_tasks"] = args.max_tasks
|
||
if args.model_dir:
|
||
config["model_dirs"] = list(dict.fromkeys(list(config.get("model_dirs") or []) + args.model_dir))
|
||
if args.no_browser:
|
||
config["open_browser"] = False
|
||
|
||
# 命令行显式指定的代码目录优先级最高，单独存一个键，避免被 config.json 里的旧值盖住
|
||
config["_cli_node_dir"] = args.node_dir
|
||
config["_path"] = str(config_path)
|
||
config["_first_run"] = not config_path.is_file()
|
||
return config
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 主流程
|
||
# --------------------------------------------------------------------------- #
|
||
def build_app(config: Dict[str, object]):
|
||
from birefnet_web import compat
|
||
from birefnet_web.engine import InferenceEngine, ModelRegistry
|
||
from birefnet_web.server import TaskManager, WebUIServer
|
||
|
||
# 1) 模型代码目录：命令行 > 环境变量 > 项目内 vendor/ > 旧配置 > 外部 ComfyUI
|
||
node_dir = compat.find_node_dir(config.get("_cli_node_dir"), config.get("node_dir") or None)
|
||
|
||
# 2) 权重目录：项目内 models/ 优先
|
||
model_dirs: List[Path] = []
|
||
for d in compat.default_model_dirs(node_dir):
|
||
if d not in model_dirs:
|
||
model_dirs.append(d)
|
||
for item in config.get("model_dirs") or []:
|
||
if item:
|
||
p = Path(str(item)).expanduser()
|
||
if p not in model_dirs:
|
||
model_dirs.append(p)
|
||
|
||
registry = ModelRegistry(model_dirs)
|
||
models = registry.scan()
|
||
|
||
output_root = Path(str(config.get("output_dir") or "outputs")).expanduser()
|
||
if not output_root.is_absolute():
|
||
output_root = PROJECT_ROOT / output_root
|
||
|
||
engine = InferenceEngine(registry, node_dir=node_dir, max_cached=int(config.get("max_cached_models") or 1))
|
||
logger = lambda msg: print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True) # noqa: E731
|
||
tasks = TaskManager(engine, output_root, max_tasks=int(config.get("max_tasks") or 50), logger=logger)
|
||
|
||
server = WebUIServer(
|
||
(str(config.get("host") or "127.0.0.1"), int(config.get("port") or 7861)),
|
||
registry=registry,
|
||
engine=engine,
|
||
tasks=tasks,
|
||
web_dir=PROJECT_ROOT / "web",
|
||
output_root=output_root,
|
||
node_dir=node_dir,
|
||
logger=logger,
|
||
project_root=PROJECT_ROOT,
|
||
)
|
||
return server, registry, models, node_dir, model_dirs, output_root
|
||
|
||
|
||
def _rel_or_abs(path: Path) -> str:
|
||
"""项目内的路径显示成相对路径，读起来更舒服。"""
|
||
try:
|
||
return str(path.resolve().relative_to(PROJECT_ROOT))
|
||
except (ValueError, OSError):
|
||
return str(path)
|
||
|
||
|
||
def print_banner(server, models, node_dir, model_dirs, output_root, config) -> None:
|
||
host, port = server.server_address[0], server.server_address[1]
|
||
shown = "127.0.0.1" if host in ("0.0.0.0", "::") else host
|
||
url = f"http://{shown}:{port}/"
|
||
from birefnet_web.engine import describe_environment
|
||
|
||
env = describe_environment(server.engine)
|
||
line = "─" * 66
|
||
print(line)
|
||
print(" BiRefNet WebUI · 智能抠图")
|
||
print(line)
|
||
print(f" 访问地址 {url}")
|
||
if host in ("0.0.0.0", "::"):
|
||
print(f" 局域网 同一网络下用 http://<本机IP>:{port}/ 访问")
|
||
print(f" 运行时 {_rel_or_abs(Path(sys.executable))} (python {env['python']})")
|
||
print(f" 推理库 torch {env['torch']} / torchvision {env['torchvision']}")
|
||
if env["cuda"]:
|
||
for dev in env["devices"]:
|
||
print(f" GPU [{dev['index']}] {dev['name']} {dev['free_mem']}G 空闲 / {dev['total_mem']}G")
|
||
else:
|
||
print(" GPU 未检测到 CUDA，将使用 CPU（较慢）")
|
||
print(f" 模型代码 {_rel_or_abs(node_dir)}")
|
||
for d in model_dirs:
|
||
mark = "✔" if Path(d).is_dir() else "✘ 不存在"
|
||
print(f" 模型目录 {_rel_or_abs(Path(d))} {mark}")
|
||
print(f" 可用模型 {len(models)} 个" + (f"：{', '.join(m.name for m in models)}" if models else "（未找到权重文件）"))
|
||
if not models:
|
||
print(" ↳ 请把 *.safetensors 放到上面的模型目录")
|
||
print(f" 输出目录 {_rel_or_abs(output_root)}")
|
||
print(f" 配置文件 {config.get('_path')}")
|
||
print(line)
|
||
print(" 按 Ctrl+C 退出")
|
||
print(line, flush=True)
|
||
|
||
|
||
def main(argv: Optional[List[str]] = None) -> int:
|
||
args = parse_args(argv)
|
||
preflight()
|
||
config = merge_config(args)
|
||
|
||
try:
|
||
server, registry, models, node_dir, model_dirs, output_root = build_app(config)
|
||
except Exception as exc:
|
||
print(f"[启动失败] {type(exc).__name__}: {exc}", file=sys.stderr)
|
||
return 1
|
||
|
||
# 首启或显式要求时落盘配置（补全探测到的路径，便于下次直接改文件）
|
||
config_to_save = {k: v for k, v in config.items() if not k.startswith("_")}
|
||
config_to_save["node_dir"] = str(node_dir)
|
||
if config.get("_first_run") or args.save_config:
|
||
save_config(Path(str(config["_path"])), config_to_save)
|
||
|
||
if args.check:
|
||
print_banner(server, models, node_dir, model_dirs, output_root, config)
|
||
print("自检完成（--check 模式，不启动服务）")
|
||
return 0
|
||
|
||
print_banner(server, models, node_dir, model_dirs, output_root, config)
|
||
|
||
if config.get("open_browser"):
|
||
host, port = server.server_address[0], server.server_address[1]
|
||
shown = "127.0.0.1" if host in ("0.0.0.0", "::") else host
|
||
threading.Thread(
|
||
target=lambda: (time.sleep(1.0), webbrowser.open(f"http://{shown}:{port}/")),
|
||
daemon=True,
|
||
).start()
|
||
|
||
try:
|
||
server.serve_forever(poll_interval=0.4)
|
||
except KeyboardInterrupt:
|
||
print("\n收到 Ctrl+C，正在退出…")
|
||
finally:
|
||
server.engine.unload()
|
||
server.server_close()
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|