rui40000andClaude Opus 4.8
5d3834503e
【修复】角色被格线切断（脚、手杖、飘起的斗篷被削掉） 新增 expand_beyond_cell（默认开启）：格子只用来判定「这是哪个方向」， 角色的实际范围由它自身的连通区域决定，按质心归属确保邻居不混入。 实测 8/8 方向的裁剪框边缘 alpha 从 1.00（内容顶到边界=被切断） 降到 0.00，S 方向高度 326→356、E 方向宽度 150→188 把缺的部分找了回来。 代价是需要两遍扫描（先求全序列并集框再提取），耗时 4.7s→14.9s。 【规则】每个参数都必须有中文 tooltip，作为以后的统一约定 全仓库 26 个节点 169 个参数，此前缺 115 个，现已 100% 覆盖。 tooltip 写「怎么调」而不只是「是什么」：给取值区间的实际影响、 推荐值与踩坑提示（如 OpenAI/ZenMux 的地址栏不能带 :// ， 素材拆分节点用于动画序列时顺序会漂移等）。 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
288 lines
12 KiB
Python
288 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Lucida 抠图节点（Ruinode）
|
||
==========================
|
||
模型：https://huggingface.co/egeorcun/lucida （MIT）
|
||
项目：https://github.com/egeorcun/lucida
|
||
|
||
Lucida 是 BiRefNet_HR 的微调版，训练目标是攻克多数开源抠图模型的短板：
|
||
伪装物体、透明材质（玻璃）、文字与 Logo、VFX 光效、插画。
|
||
作者在 203 图 9 类别基准上的 MAE（越低越好）：
|
||
伪装 0.0270、插画 0.0092、文字/Logo 0.0091（商业参考为 0.0123）、
|
||
印刷设计 0.0235、总体 0.0257。
|
||
|
||
与本仓库另两个抠图节点的分工：
|
||
- Lucida ：全自动。文字/Logo、插画、玻璃、伪装物体场景下最强
|
||
- FeyNobg：全自动。常规主体照片，速度快
|
||
- SDMatte：需框/掩码提示，适合画面中多个主体、只抠其中一个
|
||
|
||
实现说明：
|
||
- 模型代码内嵌在 lucida/ 子包（见其 __init__.py），不使用 trust_remote_code，
|
||
避免运行时联网执行远程代码；
|
||
- 预处理与 FeyNobg 完全一致（1024 双线性 + ImageNet 标准化），
|
||
直接复用那份已验证过的实现，不重复造轮子；
|
||
- 模型内部 Config.size=1024 且 dec_ipt_split=True，decoder 与 1024 输入绑定，
|
||
因此不开放分辨率选项（实测非 1024 会因通道数不匹配报错）。
|
||
"""
|
||
import os
|
||
|
||
import numpy as np
|
||
import torch
|
||
|
||
import folder_paths
|
||
|
||
# ---------------------------------------------------------------- 模型目录注册
|
||
|
||
LUCIDA_DIR = os.path.join(folder_paths.models_dir, "lucida")
|
||
os.makedirs(LUCIDA_DIR, exist_ok=True)
|
||
|
||
if "lucida" in folder_paths.folder_names_and_paths:
|
||
_paths, _exts = folder_paths.folder_names_and_paths["lucida"]
|
||
if LUCIDA_DIR not in _paths:
|
||
_paths.append(LUCIDA_DIR)
|
||
_exts.update({".safetensors", ".pth", ".pt"})
|
||
else:
|
||
folder_paths.folder_names_and_paths["lucida"] = (
|
||
[LUCIDA_DIR], {".safetensors", ".pth", ".pt"})
|
||
|
||
HF_REPO = "egeorcun/lucida"
|
||
DEFAULT_CKPT = "lucida.safetensors"
|
||
_NO_MODEL_HINT = "未找到权重，将自动下载"
|
||
|
||
_MODEL_CACHE = {}
|
||
|
||
|
||
def _list_ckpts():
|
||
files = []
|
||
if os.path.isdir(LUCIDA_DIR):
|
||
files = sorted(f for f in os.listdir(LUCIDA_DIR)
|
||
if f.lower().endswith((".safetensors", ".pth", ".pt")))
|
||
return files or [_NO_MODEL_HINT]
|
||
|
||
|
||
def _ensure_ckpt(name):
|
||
"""返回权重路径；没有就从 HuggingFace 拉一份（约 885MB）。"""
|
||
if name and name != _NO_MODEL_HINT:
|
||
p = os.path.join(LUCIDA_DIR, name)
|
||
if os.path.isfile(p):
|
||
return p
|
||
|
||
target = os.path.join(LUCIDA_DIR, DEFAULT_CKPT)
|
||
if os.path.isfile(target):
|
||
return target
|
||
|
||
print(f"[Ruinode-Lucida] 本地未找到权重，开始从 HuggingFace 下载 {HF_REPO}"
|
||
f"（约 885MB）→ {target}")
|
||
try:
|
||
from huggingface_hub import hf_hub_download
|
||
|
||
got = hf_hub_download(HF_REPO, "model.safetensors", local_dir=LUCIDA_DIR)
|
||
src = os.path.join(LUCIDA_DIR, "model.safetensors")
|
||
if os.path.isfile(src) and not os.path.isfile(target):
|
||
os.rename(src, target)
|
||
elif os.path.isfile(got) and not os.path.isfile(target):
|
||
target = got
|
||
except Exception as e:
|
||
raise RuntimeError(
|
||
f"自动下载失败：{e}\n"
|
||
f"请手动下载 https://huggingface.co/{HF_REPO} 的 model.safetensors，"
|
||
f"重命名为 {DEFAULT_CKPT} 放到：{LUCIDA_DIR}"
|
||
)
|
||
if not os.path.isfile(target):
|
||
raise RuntimeError(f"下载完成但未找到权重文件：{target}")
|
||
print(f"[Ruinode-Lucida] 下载完成：{target}")
|
||
return target
|
||
|
||
|
||
def _load_weights_strict(model, ckpt_path):
|
||
"""加载权重并严格校验，绝不接受"能跑但结果劣化"的静默失配。"""
|
||
if ckpt_path.lower().endswith(".safetensors"):
|
||
from safetensors.torch import load_file
|
||
sd = load_file(ckpt_path)
|
||
else:
|
||
sd = torch.load(ckpt_path, map_location="cpu")
|
||
sd = sd.get("state_dict", sd) if isinstance(sd, dict) else sd
|
||
|
||
# transformers 保存时可能带 model. 前缀，去掉后再比对
|
||
own = set(model.state_dict().keys())
|
||
if len(own & set(sd.keys())) == 0:
|
||
stripped = {k.split("model.", 1)[-1] if k.startswith("model.") else k: v
|
||
for k, v in sd.items()}
|
||
if len(own & set(stripped.keys())) > 0:
|
||
print("[Ruinode-Lucida] 权重键名去掉 'model.' 前缀后对齐")
|
||
sd = stripped
|
||
|
||
result = model.load_state_dict(sd, strict=False)
|
||
# BiRefNet 的 Swin 会注册 relative_position_index / attn_mask 一类
|
||
# 由窗口尺寸推出的确定性 buffer，构造时已算好，不需要也不该来自权重
|
||
benign = ("relative_position_index", "attn_mask")
|
||
bad_missing = [k for k in result.missing_keys if not k.endswith(benign)]
|
||
bad_unexpected = list(result.unexpected_keys)
|
||
if bad_missing or bad_unexpected:
|
||
for k in bad_missing[:8]:
|
||
print(f" 未被覆盖: {k}")
|
||
for k in bad_unexpected[:8]:
|
||
print(f" 未被使用: {k}")
|
||
raise RuntimeError(
|
||
f"权重与网络结构不匹配（异常缺失 {len(bad_missing)}，异常多余 "
|
||
f"{len(bad_unexpected)}）。继续推理只会得到劣化结果，故中止。\n"
|
||
f"请确认权重确为 egeorcun/lucida 的 model.safetensors。"
|
||
)
|
||
print(f"[Ruinode-Lucida] 权重加载完成（{len(own) - len(result.missing_keys)}"
|
||
f"/{len(own)} 个张量，另 {len(result.missing_keys)} 个为确定性 buffer）")
|
||
|
||
|
||
def _load_model(ckpt_path, dtype, device):
|
||
key = (ckpt_path, str(dtype), str(device))
|
||
cached = _MODEL_CACHE.get(key)
|
||
if cached is not None:
|
||
return cached
|
||
|
||
from .lucida import BiRefNet, BiRefNetConfig
|
||
|
||
print(f"[Ruinode-Lucida] 加载模型：{os.path.basename(ckpt_path)}"
|
||
f"（{dtype}，{device}）")
|
||
_MODEL_CACHE.clear() # 单份 880MB 起，不做多份缓存
|
||
# bb_pretrained=False：只搭骨架，不去联网拉 Swin 的 ImageNet 预训练权重
|
||
model = BiRefNet(config=BiRefNetConfig(bb_pretrained=False))
|
||
_load_weights_strict(model, ckpt_path)
|
||
model.eval()
|
||
model.to(device=device, dtype=dtype)
|
||
_MODEL_CACHE[key] = model
|
||
return model
|
||
|
||
|
||
class RuiLucida:
|
||
"""Lucida 全自动抠图：擅长文字/Logo、插画、玻璃与伪装物体。"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE", {
|
||
"tooltip": "待抠图的图像，支持批量（逐张推理，避免显存峰值过高）。\n"
|
||
"RGBA 输入会自动丢弃 alpha 只取 RGB。"
|
||
}),
|
||
"model_name": (_list_ckpts(), {
|
||
"tooltip": "放在 models/lucida 下的权重（.safetensors）。\n"
|
||
"未找到时首次运行会自动从 HuggingFace 下载"
|
||
"egeorcun/lucida（约 885MB）。"
|
||
}),
|
||
"precision": (["fp32", "fp16"], {
|
||
"default": "fp16",
|
||
"tooltip": "fp16 显存减半、速度更快，实测与 fp32 输出基本一致。\n"
|
||
"遇到黑图或异常时改回 fp32。"
|
||
}),
|
||
"device": (["auto", "cpu"], {
|
||
"default": "auto",
|
||
"tooltip": "auto = 有显卡就用显卡。\n"
|
||
"cpu 会慢十倍以上，仅在显存实在不够时用。"
|
||
}),
|
||
},
|
||
"optional": {
|
||
"alpha_threshold": ("FLOAT", {
|
||
"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
|
||
"tooltip": "多大置信度才算前景。模型对拿不准的区域会输出 0.5 上下的\n"
|
||
"中间值，表现为「整片主体半透明发灰」，调低可拉回不透明。\n"
|
||
"只对已有一定响应的区域有效。"
|
||
}),
|
||
"alpha_softness": ("FLOAT", {
|
||
"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01,
|
||
"tooltip": "阈值两侧过渡带宽度，决定边缘软硬。\n"
|
||
"1.0 = 完全不处理，原样输出模型结果（默认）\n"
|
||
"0.2~0.4 = 压掉灰雾但保留发丝级过渡\n"
|
||
"0 = 硬二值化。Lucida 强项之一是玻璃等真实半透明，\n"
|
||
"调小会把这类效果一并压掉，务必按素材取舍。"
|
||
}),
|
||
"keep_aspect_ratio": ("BOOLEAN", {
|
||
"default": False,
|
||
"tooltip": "模型固定吃 1024×1024，默认把图直接拉伸成正方形\n"
|
||
"（与官方用法一致）。长图/宽图形变严重时可开启，\n"
|
||
"改为等比缩放 + 边缘延展补边，推理后裁掉补边。\n"
|
||
"属试验性选项，常规比例建议关闭。"
|
||
}),
|
||
"invert_mask": ("BOOLEAN", {
|
||
"default": False,
|
||
"tooltip": "反转 alpha：默认前景为白（1），开启后前景为黑。"
|
||
}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("MASK", "IMAGE")
|
||
RETURN_NAMES = ("alpha", "cutout")
|
||
FUNCTION = "matting"
|
||
CATEGORY = "Rui-Node🐶/抠图✂️"
|
||
|
||
@classmethod
|
||
def VALIDATE_INPUTS(cls, model_name, **kwargs):
|
||
"""权重列表随目录变化，宽松放行，运行时兜底（含自动下载）。"""
|
||
return True
|
||
|
||
@staticmethod
|
||
def _remap_alpha(alpha, threshold, softness):
|
||
"""按阈值/柔和度给遮罩拉一次色阶；默认 (0.5, 1.0) 时原样返回。"""
|
||
low = threshold - softness / 2.0
|
||
high = threshold + softness / 2.0
|
||
if high <= low:
|
||
return (alpha >= threshold).to(alpha.dtype)
|
||
if abs(low) < 1e-6 and abs(high - 1.0) < 1e-6:
|
||
return alpha
|
||
return ((alpha - low) / (high - low)).clamp(0.0, 1.0)
|
||
|
||
def matting(self, image, model_name, precision, device,
|
||
alpha_threshold=0.5, alpha_softness=1.0,
|
||
keep_aspect_ratio=False, invert_mask=False):
|
||
import comfy.model_management
|
||
|
||
# BiRefNet 系模型的预处理规格相同（1024 + ImageNet 标准化），
|
||
# 直接复用 FeyNobg 节点里那份已验证的实现
|
||
from .feynobg import BiRefNetImageProcessor
|
||
|
||
ckpt = _ensure_ckpt(model_name)
|
||
dev = torch.device("cpu") if device == "cpu" \
|
||
else comfy.model_management.get_torch_device()
|
||
dtype = torch.float32 if precision == "fp32" else torch.float16
|
||
model = _load_model(ckpt, dtype, dev)
|
||
|
||
proc = BiRefNetImageProcessor(size={"height": 1024, "width": 1024})
|
||
|
||
B, H, W, C = image.shape
|
||
if C == 4:
|
||
image = image[..., :3]
|
||
elif C == 1:
|
||
image = image.repeat(1, 1, 1, 3)
|
||
elif C != 3:
|
||
raise ValueError(f"image 需要 1 / 3 / 4 通道，实际收到 {C} 通道")
|
||
image = image.contiguous()
|
||
|
||
alphas = []
|
||
for i in range(B):
|
||
# 逐张推理：Swin-Large 在 1024 下峰值显存不低，整批一次容易 OOM
|
||
pixel_values, meta = proc.preprocess_tensor(
|
||
image[i:i + 1], device=dev, dtype=dtype,
|
||
keep_aspect=bool(keep_aspect_ratio))
|
||
with torch.no_grad():
|
||
preds = model(pixel_values)
|
||
# eval 模式下 BiRefNet.forward 返回各尺度预测的列表，最后一项是最终结果
|
||
logits = preds[-1] if isinstance(preds, (list, tuple)) else preds
|
||
alpha = proc.post_process_alpha_matting(
|
||
{"logits": logits.float()}, target_sizes=[(H, W)], crop=meta)[0]
|
||
alphas.append(alpha.clamp(0, 1).cpu())
|
||
|
||
alpha = torch.stack(alphas) # [B,H,W]
|
||
alpha = self._remap_alpha(alpha, float(alpha_threshold),
|
||
float(alpha_softness))
|
||
if invert_mask:
|
||
alpha = 1.0 - alpha
|
||
|
||
cutout = image.detach().cpu().float() * alpha.unsqueeze(-1)
|
||
return (alpha, cutout)
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"RuiLucida": RuiLucida,
|
||
}
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"RuiLucida": "Lucida 抠图 / Lucida Matting",
|
||
}
|