500 lines
19 KiB
Python
500 lines
19 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""模型注册表与推理引擎。
|
||
|
||
设计要点：
|
||
* 模型代码直接复用 ComfyUI 节点（见 compat.py），不复制、不魔改；
|
||
* 模型按 (路径, mtime, 设备, 精度, 架构) 做 LRU 缓存，默认只驻留 1 个，避免显存爆炸；
|
||
* float16 / bfloat16 推理若出现 NaN/Inf 会自动回退 float32 重算一次并记录告警。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import gc
|
||
import os
|
||
import threading
|
||
import time
|
||
from collections import OrderedDict
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
from typing import Callable, Dict, List, Optional, Tuple
|
||
|
||
import numpy as np
|
||
import torch
|
||
from PIL import Image
|
||
|
||
from . import compat, imageops
|
||
|
||
ProgressFn = Optional[Callable[[str, float], None]]
|
||
|
||
#: 不是抠图模型，而是骨干网络权重，扫描时跳过
|
||
_BACKBONE_PREFIXES = ("swin_", "pvt_v2_")
|
||
|
||
#: 支持加载的权重后缀
|
||
_MODEL_SUFFIXES = (".safetensors", ".pth", ".pt", ".ckpt")
|
||
|
||
#: 精度选项 -> torch dtype；auto 由设备决定
|
||
DTYPES: Dict[str, Optional[torch.dtype]] = {
|
||
"auto": None,
|
||
"float32": torch.float32,
|
||
"float16": torch.float16,
|
||
"bfloat16": torch.bfloat16,
|
||
}
|
||
|
||
#: 前景精修超过该像素量时自动跳过（避免内存与耗时失控）
|
||
_REFINE_PIXEL_LIMIT = 40_000_000
|
||
|
||
|
||
@dataclass
|
||
class ModelInfo:
|
||
"""一个可用的模型权重文件。"""
|
||
|
||
key: str
|
||
name: str
|
||
file: str
|
||
path: str
|
||
size_mb: float
|
||
mtime: float
|
||
arch: str # "v1" | "old"
|
||
bb_index: int
|
||
backbone: str
|
||
directory: str
|
||
|
||
def to_dict(self) -> dict:
|
||
return {
|
||
"key": self.key,
|
||
"name": self.name,
|
||
"file": self.file,
|
||
"path": self.path,
|
||
"size_mb": round(self.size_mb, 1),
|
||
"arch": self.arch,
|
||
"bb_index": self.bb_index,
|
||
"backbone": self.backbone,
|
||
"directory": self.directory,
|
||
}
|
||
|
||
|
||
def guess_bb_index(filename: str) -> int:
|
||
"""根据文件名猜骨干网络：lite 系列用 swin_v1_t(3)，其余用 swin_v1_l(6)。"""
|
||
return 3 if "lite" in filename.lower() else 6
|
||
|
||
|
||
def guess_arch(filename: str) -> str:
|
||
return "old" if os.path.basename(filename) in compat.OLD_MODEL_FILES else "v1"
|
||
|
||
|
||
class ModelRegistry:
|
||
"""扫描并维护可用模型列表。"""
|
||
|
||
def __init__(self, model_dirs: List[Path]) -> None:
|
||
self.model_dirs = [Path(d) for d in model_dirs]
|
||
self._models: "OrderedDict[str, ModelInfo]" = OrderedDict()
|
||
self._lock = threading.Lock()
|
||
|
||
def scan(self) -> List[ModelInfo]:
|
||
found: "OrderedDict[str, ModelInfo]" = OrderedDict()
|
||
for idx, d in enumerate(self.model_dirs):
|
||
if not d or not d.is_dir():
|
||
continue
|
||
try:
|
||
entries = sorted(d.iterdir(), key=lambda p: p.name.lower())
|
||
except OSError:
|
||
continue
|
||
for entry in entries:
|
||
if not entry.is_file():
|
||
continue
|
||
if entry.suffix.lower() not in _MODEL_SUFFIXES:
|
||
continue
|
||
if entry.name.lower().startswith(_BACKBONE_PREFIXES):
|
||
continue
|
||
try:
|
||
stat = entry.stat()
|
||
except OSError:
|
||
continue
|
||
stem = entry.stem
|
||
bb_index = guess_bb_index(stem)
|
||
bb = {3: "swin_v1_t", 6: "swin_v1_l"}[bb_index]
|
||
key = f"{stem}@{idx}" if stem in found else stem
|
||
if key in found: # 极少数重名情况
|
||
key = f"{stem}@{idx}"
|
||
found[key] = ModelInfo(
|
||
key=key,
|
||
name=stem,
|
||
file=entry.name,
|
||
path=str(entry.resolve()),
|
||
size_mb=stat.st_size / (1024 * 1024),
|
||
mtime=stat.st_mtime,
|
||
arch=guess_arch(entry.name),
|
||
bb_index=bb_index,
|
||
backbone=bb,
|
||
directory=str(d),
|
||
)
|
||
with self._lock:
|
||
self._models = found
|
||
return list(found.values())
|
||
|
||
def list(self) -> List[ModelInfo]:
|
||
with self._lock:
|
||
if not self._models:
|
||
return self.scan()
|
||
return list(self._models.values())
|
||
|
||
def get(self, key: str) -> ModelInfo:
|
||
for info in self.list():
|
||
if key in (info.key, info.name, info.file, info.path):
|
||
return info
|
||
raise KeyError(f"未找到模型 {key!r}，可用模型：{[m.key for m in self.list()]}")
|
||
|
||
|
||
class InferenceEngine:
|
||
"""串行执行抠图推理（GPU 上同一时刻只跑一个模型）。"""
|
||
|
||
def __init__(self, registry: ModelRegistry, node_dir: Path, max_cached: int = 1) -> None:
|
||
self.registry = registry
|
||
self.node_dir = Path(node_dir)
|
||
self.max_cached = max(1, int(max_cached))
|
||
self._cache: "OrderedDict[tuple, Tuple[object, str, torch.dtype]]" = OrderedDict()
|
||
self._lock = threading.RLock()
|
||
self._classes: Optional[dict] = None
|
||
self.last_warning: Optional[str] = None
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# 模型加载
|
||
# ------------------------------------------------------------------ #
|
||
def _model_classes(self) -> dict:
|
||
if self._classes is None:
|
||
self._classes = compat.load_model_classes(self.node_dir, self.registry.model_dirs)
|
||
return self._classes
|
||
|
||
@staticmethod
|
||
def resolve_device(device: str) -> str:
|
||
device = (device or "auto").strip().lower()
|
||
if device in ("auto", "", "gpu"):
|
||
return "cuda" if torch.cuda.is_available() else "cpu"
|
||
if device.startswith("cuda"):
|
||
if not torch.cuda.is_available():
|
||
raise RuntimeError("请求使用 CUDA，但当前 torch 未检测到可用 GPU")
|
||
return device
|
||
return "cpu"
|
||
|
||
@staticmethod
|
||
def resolve_dtype(dtype: str, device: str) -> torch.dtype:
|
||
dt = DTYPES.get((dtype or "auto").lower(), None)
|
||
if dt is not None:
|
||
return dt
|
||
# auto：CUDA 下用 fp32 权重 + fp16 autocast（官方推荐），CPU 下 fp32
|
||
return torch.float32
|
||
|
||
@staticmethod
|
||
def _autocast_dtype(dtype: str, device: str) -> Optional[torch.dtype]:
|
||
if not device.startswith("cuda"):
|
||
return None
|
||
key = (dtype or "auto").lower()
|
||
if key in ("float16", "auto"):
|
||
return torch.float16
|
||
if key == "bfloat16":
|
||
return torch.bfloat16
|
||
return None
|
||
|
||
def _read_state_dict(self, info: ModelInfo) -> dict:
|
||
if info.path.lower().endswith(".safetensors"):
|
||
import safetensors.torch
|
||
|
||
return safetensors.torch.load_file(info.path, device="cpu")
|
||
try:
|
||
sd = torch.load(info.path, map_location="cpu", weights_only=True)
|
||
except Exception:
|
||
sd = torch.load(info.path, map_location="cpu", weights_only=False)
|
||
check_state_dict = self._model_classes()["check_state_dict"]
|
||
return check_state_dict(sd)
|
||
|
||
def _build_model(self, info: ModelInfo, state_dict: dict, dtype: torch.dtype):
|
||
"""构建网络并载入权重；骨干索引猜错时自动换一个重试。"""
|
||
classes = self._model_classes()
|
||
BiRefNet = classes["BiRefNet"]
|
||
OldBiRefNet = classes["OldBiRefNet"]
|
||
|
||
if info.arch == "old":
|
||
if OldBiRefNet is None:
|
||
raise RuntimeError("该权重属于旧版架构，但节点内缺少 birefnet_old 包")
|
||
candidates = [-1]
|
||
else:
|
||
candidates = [info.bb_index] + [i for i in (6, 3) if i != info.bb_index]
|
||
|
||
last_err: Optional[Exception] = None
|
||
for bb_index in candidates:
|
||
try:
|
||
if bb_index < 0:
|
||
model = OldBiRefNet(bb_pretrained=False)
|
||
version = "old"
|
||
else:
|
||
model = BiRefNet(bb_pretrained=False, bb_index=bb_index)
|
||
version = "v1"
|
||
if dtype != torch.float32:
|
||
model = model.to(dtype=dtype)
|
||
model.load_state_dict(state_dict)
|
||
if bb_index >= 0 and bb_index != info.bb_index:
|
||
self.last_warning = (
|
||
f"模型 {info.name} 的骨干索引自动修正为 {bb_index}（文件名推断为 {info.bb_index}）"
|
||
)
|
||
return model, version, bb_index
|
||
except Exception as exc: # 尺寸不匹配等
|
||
last_err = exc
|
||
continue
|
||
raise RuntimeError(f"加载模型 {info.name} 失败：{last_err}") from last_err
|
||
|
||
def get_model(self, key: str, device: str = "auto", dtype: str = "auto", arch: str = "auto"):
|
||
info = self.registry.get(key)
|
||
if arch in ("v1", "old"):
|
||
info = ModelInfo(**{**info.__dict__, "arch": arch})
|
||
device = self.resolve_device(device)
|
||
torch_dtype = self.resolve_dtype(dtype, device)
|
||
cache_key = (info.path, info.mtime, device, str(torch_dtype), info.arch)
|
||
|
||
with self._lock:
|
||
if cache_key in self._cache:
|
||
self._cache.move_to_end(cache_key)
|
||
model, version, _ = self._cache[cache_key]
|
||
return model, version, info
|
||
state_dict = self._read_state_dict(info)
|
||
model, version, _ = self._build_model(info, state_dict, torch_dtype)
|
||
del state_dict
|
||
model = model.to(device)
|
||
model.eval()
|
||
for p in model.parameters():
|
||
p.requires_grad_(False)
|
||
self._cache[cache_key] = (model, version, torch_dtype)
|
||
self._cache.move_to_end(cache_key)
|
||
self._evict_locked()
|
||
return model, version, info
|
||
|
||
def _evict_locked(self) -> None:
|
||
while len(self._cache) > self.max_cached:
|
||
_, (model, _, _) = self._cache.popitem(last=False)
|
||
del model
|
||
gc.collect()
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
|
||
def unload(self) -> None:
|
||
with self._lock:
|
||
self._cache.clear()
|
||
gc.collect()
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# 推理
|
||
# ------------------------------------------------------------------ #
|
||
def remove_background(
|
||
self,
|
||
image: Image.Image,
|
||
options: Dict[str, object],
|
||
progress: ProgressFn = None,
|
||
) -> Dict[str, object]:
|
||
"""对单张图片执行抠图。
|
||
|
||
options 支持（均有默认值）：
|
||
model / device / dtype / arch
|
||
resolution_mode: square | longest | custom
|
||
width / height / longest_side
|
||
upscale_method / mask_threshold
|
||
refine_foreground / blur_size / blur_size_two
|
||
background: transparent | color
|
||
bg_color: "#ffffff"
|
||
output_mask: bool（默认 False：不额外产出遮罩文件）
|
||
final_longest_side: int (0 = 保持原图尺寸；>0 = 等比缩放到该最长边)
|
||
"""
|
||
t0 = time.perf_counter()
|
||
self.last_warning = None
|
||
|
||
def report(stage: str, frac: float) -> None:
|
||
if progress:
|
||
try:
|
||
progress(stage, max(0.0, min(1.0, frac)))
|
||
except Exception:
|
||
pass
|
||
|
||
model_key = str(options.get("model") or "")
|
||
if not model_key:
|
||
raise ValueError("未指定模型")
|
||
device = str(options.get("device") or "auto")
|
||
dtype = str(options.get("dtype") or "auto")
|
||
arch = str(options.get("arch") or "auto")
|
||
|
||
report("加载模型", 0.05)
|
||
model, version, info = self.get_model(model_key, device, dtype, arch)
|
||
model_dtype = next(model.parameters()).dtype
|
||
model_device = next(model.parameters()).device
|
||
|
||
report("预处理", 0.2)
|
||
rgb = imageops.pil_to_rgb_array(image)
|
||
src_h, src_w = rgb.shape[:2]
|
||
in_h, in_w = imageops.build_input_size(
|
||
src_h,
|
||
src_w,
|
||
mode=str(options.get("resolution_mode") or "square"),
|
||
width=int(options.get("width") or 1024),
|
||
height=int(options.get("height") or 1024),
|
||
longest_side=int(options.get("longest_side") or 1024),
|
||
)
|
||
upscale_method = str(options.get("upscale_method") or "bilinear")
|
||
tensor = imageops.preprocess(rgb, (in_h, in_w), upscale_method)
|
||
|
||
report("推理", 0.35)
|
||
x = tensor.to(model_device)
|
||
if model_dtype != torch.float32:
|
||
x = x.to(model_dtype)
|
||
|
||
autocast_dtype = self._autocast_dtype(dtype, device)
|
||
alpha = self._forward(model, x, model_device, model_dtype, autocast_dtype)
|
||
del x
|
||
|
||
# 精度兜底：半精度偶发 NaN 时用 fp32 重算一次
|
||
if not bool(torch.isfinite(alpha).all()):
|
||
self.last_warning = "半精度推理出现 NaN/Inf，已自动回退 float32 重算"
|
||
self.unload()
|
||
model, version, info = self.get_model(model_key, device, "float32", arch)
|
||
model_device = next(model.parameters()).device
|
||
x = tensor.to(model_device)
|
||
alpha = self._forward(model, x, model_device, torch.float32, None)
|
||
del x
|
||
del tensor
|
||
|
||
report("后处理", 0.75)
|
||
alpha = imageops.upscale_mask(alpha, src_h, src_w, upscale_method)
|
||
alpha = imageops.filter_mask(alpha, float(options.get("mask_threshold") or 0.0))
|
||
alpha_np = alpha.squeeze(0).squeeze(0).to(torch.float32).cpu().numpy()
|
||
del alpha
|
||
alpha_np = np.clip(alpha_np, 0.0, 1.0)
|
||
|
||
want_mask = bool(options.get("output_mask", False))
|
||
background = str(options.get("background") or "transparent")
|
||
rgb_out = rgb
|
||
refine = bool(options.get("refine_foreground", True))
|
||
if refine and src_h * src_w > _REFINE_PIXEL_LIMIT:
|
||
refine = False
|
||
self.last_warning = (
|
||
f"图像像素 {src_h * src_w / 1e6:.1f}MP 超过前景精修上限，已自动跳过该步骤"
|
||
)
|
||
|
||
if refine:
|
||
report("前景精修", 0.85)
|
||
rgb_out = np.clip(
|
||
np.rint(
|
||
imageops.refine_foreground(
|
||
rgb.astype(np.float32) / 255.0,
|
||
alpha_np,
|
||
int(options.get("blur_size") or 90),
|
||
int(options.get("blur_size_two") or 6),
|
||
)
|
||
* 255.0
|
||
),
|
||
0,
|
||
255,
|
||
).astype(np.uint8)
|
||
|
||
if background == "color":
|
||
cutout = imageops.compose_on_color(
|
||
rgb_out, alpha_np, imageops.parse_color(options.get("bg_color"))
|
||
)
|
||
else:
|
||
cutout = imageops.to_pil_rgba(rgb_out, alpha_np)
|
||
|
||
mask_pil = imageops.to_pil_mask(alpha_np) if want_mask else None
|
||
|
||
# 最终尺寸：按原图比例把最长边缩放到指定像素（0 = 保持原尺寸）
|
||
final_side = int(options.get("final_longest_side") or 0)
|
||
new_size = imageops.fit_longest_side(cutout.size, final_side)
|
||
if new_size != cutout.size:
|
||
cutout = cutout.resize(new_size, Image.LANCZOS)
|
||
if mask_pil is not None:
|
||
mask_pil = mask_pil.resize(new_size, Image.LANCZOS)
|
||
|
||
report("保存结果", 0.95)
|
||
# 前景占比：便于前端提示「模型可能没抠到东西」
|
||
coverage = float((alpha_np > 0.5).mean()) if alpha_np.size else 0.0
|
||
return {
|
||
"cutout": cutout,
|
||
"mask": mask_pil,
|
||
"width": cutout.width,
|
||
"height": cutout.height,
|
||
"source_width": src_w,
|
||
"source_height": src_h,
|
||
"input_size": [in_w, in_h],
|
||
"coverage": round(coverage, 4),
|
||
"model": info.name,
|
||
"model_path": info.path,
|
||
"backbone": info.backbone,
|
||
"arch": version,
|
||
"device": str(model_device),
|
||
"dtype": str(model_dtype),
|
||
"elapsed": round(time.perf_counter() - t0, 2),
|
||
"warning": self.last_warning,
|
||
"notes": self._notes(background, refine),
|
||
}
|
||
|
||
@staticmethod
|
||
def _notes(background: str, refine: bool) -> List[str]:
|
||
notes = []
|
||
if refine:
|
||
notes.append("已启用前景精修（去除边缘残留背景色）")
|
||
if background == "color":
|
||
notes.append("已合成到自定义背景色")
|
||
return notes
|
||
|
||
@staticmethod
|
||
def _forward(model, x, device: torch.device, model_dtype: torch.dtype, autocast_dtype):
|
||
"""前向一次，返回 sigmoid 后的低分辨率遮罩 (1,1,h,w)。"""
|
||
use_autocast = autocast_dtype is not None and str(device).startswith("cuda")
|
||
ctx = (
|
||
torch.autocast(device_type="cuda", dtype=autocast_dtype, enabled=True)
|
||
if use_autocast
|
||
else torch.autocast(device_type="cpu", enabled=False)
|
||
)
|
||
with torch.inference_mode(), ctx:
|
||
out = model(x)
|
||
pred = out[-1] if isinstance(out, (list, tuple)) else out
|
||
alpha = pred.sigmoid().float()
|
||
return alpha
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 环境信息
|
||
# --------------------------------------------------------------------------- #
|
||
def describe_environment(engine: InferenceEngine) -> dict:
|
||
"""给前端展示的设备/运行环境信息。"""
|
||
try:
|
||
import torch as _torch # noqa: F401
|
||
import torchvision as _tv
|
||
|
||
tv_version = getattr(_tv, "__version__", "?")
|
||
except Exception: # pragma: no cover
|
||
tv_version = "?"
|
||
|
||
devices = []
|
||
if torch.cuda.is_available():
|
||
for i in range(torch.cuda.device_count()):
|
||
try:
|
||
props = torch.cuda.get_device_properties(i)
|
||
free, total = torch.cuda.mem_get_info(i)
|
||
devices.append(
|
||
{
|
||
"index": i,
|
||
"name": props.name,
|
||
"total_mem": round(total / 1024**3, 1),
|
||
"free_mem": round(free / 1024**3, 1),
|
||
"cc": f"{props.major}.{props.minor}",
|
||
}
|
||
)
|
||
except Exception:
|
||
devices.append({"index": i, "name": "CUDA 设备", "total_mem": 0, "free_mem": 0})
|
||
return {
|
||
"python": os.sys.version.split()[0],
|
||
"torch": torch.__version__,
|
||
"torchvision": tv_version,
|
||
"cuda": torch.cuda.is_available(),
|
||
"cuda_version": getattr(torch.version, "cuda", None),
|
||
"devices": devices,
|
||
"default_device": "cuda" if torch.cuda.is_available() else "cpu",
|
||
}
|