feat: initial project setup

2026-10-08 09:46:47 +08:00
commit 5833a303bc
65 changed files with 11628 additions and 0 deletions
@@ -0,0 +1,499 @@
+# -*- 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",
+    }