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>
376 lines
17 KiB
Python
376 lines
17 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
SDMatte 精细化抠图节点（Ruinode）
|
||
|
||
对应论文：SDMatte: Grafting Diffusion Models for Interactive Matting (ICCV 2025)
|
||
官方实现：https://github.com/vivoCameraResearch/SDMatte
|
||
|
||
与社区已有实现的关键差别（按实测影响从大到小排列）：
|
||
|
||
1. 视觉提示走官方的 bbox_mask 路径，并真正传入归一化坐标。
|
||
官方 configs/SDMatte.py 固定 aux_input="bbox_mask"，其 aux_input_list 只有
|
||
point_mask / bbox_mask / mask —— trimap 从未作为视觉提示参与训练。
|
||
ComfyUI-SDMatte 传 aux_input="trimap"，等于把模型推到没训练过的输入模式上，
|
||
且该分支的 trimap_coords 恒为 [0,0,1,1]，定位信息全部丢失。
|
||
实测（官方效果图里的羊驼，与官方公布 alpha 比）：本节点 MAD=0.0113，
|
||
ComfyUI-SDMatte MAD=0.0884，相差 7.8 倍，且其输出明显发灰、边缘晕开。
|
||
|
||
2. UNet 结构用官方 LongfeiHuang/SDMatte 的 config.json 构建，而非原版 SD 2.1 的。
|
||
官方配置额外定义了 bbox_time_embed_dim=320 等三个字段；用 SD 2.1 的配置会缺字段，
|
||
只能靠猜默认值，一旦猜错，对应权重会被 strict=False 静默丢弃。
|
||
|
||
3. 照搬官方 configs/SDMatte.py 的 model_kwargs，含
|
||
use_encoder_hidden_states_list=[False, True, False]（漏传会退化成 [True,True,True]）。
|
||
实测该项单独影响不大（羊驼 MAD 0.01135 -> 0.01148），透明物体上更明显；
|
||
影响虽小，但没有任何理由偏离官方配置。
|
||
|
||
4. 全程 fp32、1024 分辨率推理，与官方测试配置一致，且不做任何启发式后处理。
|
||
ComfyUI-SDMatte 的 mask_refine 会做阈值截断与 *1.2 提亮，实测反而把边缘打成硬边。
|
||
|
||
Stable Diffusion 2.1 的权重在本流程中不需要：官方 load_weight=False，
|
||
网络只从 config 建骨架，全部权重来自 SDMatte 检查点。
|
||
"""
|
||
|
||
import os
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import torch
|
||
|
||
import folder_paths
|
||
|
||
# ---------------------------------------------------------------- 模型目录注册
|
||
|
||
SDMATTE_DIR = os.path.join(folder_paths.models_dir, "SDMatte")
|
||
os.makedirs(SDMATTE_DIR, exist_ok=True)
|
||
|
||
# 与 ComfyUI-SDMatte 共用同一目录，已下载过的权重可直接复用
|
||
if "SDMatte" in folder_paths.folder_names_and_paths:
|
||
_paths, _exts = folder_paths.folder_names_and_paths["SDMatte"]
|
||
if SDMATTE_DIR not in _paths:
|
||
_paths.append(SDMATTE_DIR)
|
||
_exts.update({".pth", ".safetensors", ".pt", ".ckpt"})
|
||
else:
|
||
folder_paths.folder_names_and_paths["SDMatte"] = (
|
||
[SDMATTE_DIR],
|
||
{".pth", ".safetensors", ".pt", ".ckpt"},
|
||
)
|
||
|
||
# 官方配置（unet/vae/text_encoder/scheduler/tokenizer）已随节点一起分发，无需联网
|
||
CONFIG_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sdmatte", "configs")
|
||
|
||
# 官方 configs/SDMatte.py -> hy_dict.model_kwargs，逐字对齐
|
||
OFFICIAL_MODEL_KWARGS = dict(
|
||
load_weight=False,
|
||
conv_scale=3,
|
||
num_inference_steps=1,
|
||
aux_input="bbox_mask",
|
||
add_noise=False,
|
||
use_dis_loss=True,
|
||
use_aux_input=True,
|
||
use_coor_input=True,
|
||
use_attention_mask=True,
|
||
residual_connection=False,
|
||
use_encoder_hidden_states=True,
|
||
use_attention_mask_list=[True, True, True],
|
||
use_encoder_hidden_states_list=[False, True, False],
|
||
)
|
||
|
||
_MODEL_CACHE = {}
|
||
|
||
|
||
def _list_checkpoints():
|
||
try:
|
||
files = folder_paths.get_filename_list("SDMatte")
|
||
except Exception:
|
||
files = []
|
||
if not files:
|
||
files = [
|
||
f for f in os.listdir(SDMATTE_DIR)
|
||
if f.lower().endswith((".pth", ".safetensors", ".pt", ".ckpt"))
|
||
] if os.path.isdir(SDMATTE_DIR) else []
|
||
return sorted(files) if files else ["未找到权重，请放入 models/SDMatte"]
|
||
|
||
|
||
def _resolve_ckpt(name):
|
||
path = folder_paths.get_full_path("SDMatte", name)
|
||
if path and os.path.isfile(path):
|
||
return path
|
||
direct = os.path.join(SDMATTE_DIR, name)
|
||
if os.path.isfile(direct):
|
||
return direct
|
||
raise FileNotFoundError(
|
||
f"找不到权重 '{name}'。请把 SDMatte_plus.pth 放到：{SDMATTE_DIR}\n"
|
||
"官方下载地址：https://huggingface.co/LongfeiHuang/SDMatte"
|
||
)
|
||
|
||
|
||
def _build_model(ckpt_path, dtype, device, attention_slicing=True):
|
||
from .sdmatte.ckpt_io import load_sdmatte_state_dict, adapt_state_dict_to_model
|
||
from .sdmatte.meta_arch import SDMatte
|
||
|
||
print(f"[Ruinode-SDMatte] 按官方配置构建网络：{CONFIG_DIR}")
|
||
model = SDMatte(pretrained_model_name_or_path=CONFIG_DIR, **OFFICIAL_MODEL_KWARGS)
|
||
|
||
print(f"[Ruinode-SDMatte] 读取权重：{os.path.basename(ckpt_path)}")
|
||
state_dict = load_sdmatte_state_dict(ckpt_path)
|
||
print(f"[Ruinode-SDMatte] 权重张量数：{len(state_dict)}")
|
||
|
||
state_dict = adapt_state_dict_to_model(state_dict, model)
|
||
missing, unexpected = model.load_state_dict(state_dict, strict=False)
|
||
|
||
# load_state_dict(strict=False) 会把对不上的权重悄悄丢掉，模型照样能跑，
|
||
# 只是输出质量下降 —— 这是最难排查的一类问题，所以这里必须叫停而不是继续。
|
||
if missing or unexpected:
|
||
print(f"[Ruinode-SDMatte] 权重未对齐：缺失 {len(missing)} 个，多余 {len(unexpected)} 个")
|
||
for k in list(missing)[:10]:
|
||
print(f" 未被覆盖: {k}")
|
||
for k in list(unexpected)[:10]:
|
||
print(f" 未被使用: {k}")
|
||
raise RuntimeError(
|
||
f"权重与网络结构不匹配（缺失 {len(missing)}，多余 {len(unexpected)}）。"
|
||
"继续推理会得到质量劣化的结果，故中止。请确认权重文件是否为官方 SDMatte / SDMatte_plus。"
|
||
)
|
||
|
||
print(f"[Ruinode-SDMatte] 权重与网络完全匹配（{len(state_dict)} 个张量）")
|
||
|
||
# 1024 分辨率下最浅一层的自注意力是 16384x16384，一次性算完峰值约 15.5GB。
|
||
# 分片逐块计算同一批注意力，实测显存降到 9.1GB 且更快（省下的搬运多于分片开销），
|
||
# 数值仅因浮点累加次序不同产生 ~1e-6 的偏差，肉眼不可见。
|
||
if attention_slicing:
|
||
try:
|
||
from diffusers.models.attention_processor import SlicedAttnProcessor
|
||
|
||
model.unet.set_attn_processor(SlicedAttnProcessor(slice_size=1))
|
||
print("[Ruinode-SDMatte] 已启用注意力分片（显存约降 40%）")
|
||
except Exception as e:
|
||
print(f"[Ruinode-SDMatte] 注意力分片启用失败，按不分片继续：{e}")
|
||
|
||
model.eval()
|
||
model.to(device=device, dtype=dtype)
|
||
return model
|
||
|
||
|
||
class RuiSDMatteLoader:
|
||
"""加载 SDMatte 权重。支持官方 .pth（12.1GB）与社区转换的 .safetensors（5.19GB）。"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"ckpt_name": (_list_checkpoints(), {
|
||
"tooltip": "放在 models/SDMatte 下的权重。\n"
|
||
"官方 SDMatte_plus.pth 与社区 SDMatte_plus.safetensors 的模型权重完全等价，\n"
|
||
"pth 多出的约 6GB 是训练用的优化器状态，推理不参与。"
|
||
}),
|
||
"precision": (["fp32", "fp16"], {
|
||
"default": "fp32",
|
||
"tooltip": "官方测试配置为 fp32（amp.enabled=False）。\n"
|
||
"fp16 省显存但 SD 2.1 的 VAE 在半精度下容易溢出，可能出现黑图或噪点。"
|
||
}),
|
||
"device": (["auto", "cpu"], {
|
||
"default": "auto",
|
||
"tooltip": "auto = 有显卡就用显卡。\n"
|
||
"SDMatte 体量大，cpu 推理会慢到不实用，仅作兜底。"
|
||
}),
|
||
"attention_slicing": ("BOOLEAN", {
|
||
"default": True,
|
||
"tooltip": "分片计算注意力。1024 分辨率下显存峰值从约 15.5GB 降到 9.1GB，\n"
|
||
"实测速度反而略快，输出差异在 1e-6 量级、肉眼不可见。\n"
|
||
"显存充裕且想严格对齐官方数值时可关闭。"
|
||
}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("SDMATTE_MODEL",)
|
||
RETURN_NAMES = ("sdmatte_model",)
|
||
FUNCTION = "load"
|
||
CATEGORY = "Ruinode/SDMatte"
|
||
|
||
def load(self, ckpt_name, precision, device, attention_slicing=True):
|
||
import comfy.model_management
|
||
|
||
ckpt_path = _resolve_ckpt(ckpt_name)
|
||
dev = torch.device("cpu") if device == "cpu" else comfy.model_management.get_torch_device()
|
||
dtype = torch.float32 if precision == "fp32" else torch.float16
|
||
|
||
key = (ckpt_path, str(dtype), str(dev), bool(attention_slicing))
|
||
cached = _MODEL_CACHE.get(key)
|
||
if cached is not None:
|
||
return (cached,)
|
||
|
||
_MODEL_CACHE.clear() # 单份 5GB 起步，不做多份缓存
|
||
model = _build_model(ckpt_path, dtype, dev, attention_slicing)
|
||
_MODEL_CACHE[key] = model
|
||
return (model,)
|
||
|
||
|
||
class RuiSDMatte:
|
||
"""用视觉提示（框/掩码/点）驱动 SDMatte，输出精细 alpha。"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"sdmatte_model": ("SDMATTE_MODEL", {
|
||
"tooltip": "由「SDMatte 加载器」输出的模型。\n"
|
||
"加载器有缓存，多个抠图节点可共用同一个加载器。"
|
||
}),
|
||
"image": ("IMAGE", {
|
||
"tooltip": "待抠图的原图。RGBA 输入会自动丢弃 alpha 只取 RGB\n"
|
||
"（VAE 编码器只收 3 通道，多出的 alpha 不能当颜色喂进去）。"
|
||
}),
|
||
"mask": ("MASK", {
|
||
"tooltip": "指示要抠哪个目标的提示掩码，不必精确，粗略覆盖主体即可。"
|
||
}),
|
||
"prompt_type": (["bbox_mask", "mask", "point_mask", "auto_mask"], {
|
||
"default": "bbox_mask",
|
||
"tooltip": "视觉提示类型。\n"
|
||
"bbox_mask：取掩码外接框作为提示，官方测试脚本的默认路径，通常最稳；\n"
|
||
"mask：直接用掩码本身，适合已有较准的粗分割；\n"
|
||
"point_mask：在掩码内随机取 10 个点。\n"
|
||
" ⚠ 仅 SDMatte.pth 支持；SDMatte_plus.pth 用它会输出全黑，\n"
|
||
" 因为 SDMatte*(plus) 的训练集用 COCO-Matte 替换了 RefMatte，未训练点提示；\n"
|
||
"auto_mask：不给定位信息，全图自动，画面只有单一主体时可用。"
|
||
}),
|
||
"inference_size": ([512, 640, 768, 896, 1024, 1152, 1280], {
|
||
"default": 1024,
|
||
"tooltip": "官方测试固定用 1024，降低会明显损失边缘细节。"
|
||
}),
|
||
"is_transparent": ("BOOLEAN", {
|
||
"default": False,
|
||
"tooltip": "目标是否为玻璃、纱、烟雾等透明/半透明物体。\n"
|
||
"该开关会切换模型的不透明度嵌入分支，抠透明物时务必打开。"
|
||
}),
|
||
},
|
||
"optional": {
|
||
"caption": ("STRING", {
|
||
"default": "", "multiline": False,
|
||
"tooltip": "目标物体的英文描述（对应 RefMatte 的指代表达式），经 CLIP 编码后\n"
|
||
"注入 UNet 的下采样/上采样两段 cross-attention。\n"
|
||
"· SDMatte.pth：有语义作用，填对能小幅提升（实测羊驼图 MAD 0.01120→0.01072）\n"
|
||
"· SDMatte_plus.pth：无语义作用，填了反而更差，请留空\n"
|
||
" （plus 的训练集用 COCO-Matte 替换了 RefMatte，未学过文本指代）\n"
|
||
"留空即官方测试时的默认行为。"
|
||
}),
|
||
"point_radius": ("INT", {
|
||
"default": 35, "min": 5, "max": 100,
|
||
"tooltip": "仅 prompt_type=point_mask 时生效，其余模式完全不参与运算。\n"
|
||
"含义：在提示区域随机取 10 个点，每点用 sigma=该值的高斯核晕开成光斑，\n"
|
||
"再逐像素取最大值合成提示图 —— 即每个点的影响半径。\n"
|
||
"官方训练用 25，测试期用 25+10=35，故默认 35。"
|
||
}),
|
||
"seed": ("INT", {
|
||
"default": 0, "min": 0, "max": 0xFFFFFFFF,
|
||
"tooltip": "随机种子，仅在 prompt_type=point_mask 时真正起作用\n"
|
||
"（决定在提示区域里随机取哪 10 个点）。\n"
|
||
"其余提示类型下改它不会改变结果。"
|
||
}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("MASK", "IMAGE")
|
||
RETURN_NAMES = ("alpha", "cutout")
|
||
FUNCTION = "apply"
|
||
CATEGORY = "Ruinode/SDMatte"
|
||
|
||
def apply(self, sdmatte_model, image, mask, prompt_type, inference_size,
|
||
is_transparent, caption="", point_radius=35, seed=0):
|
||
from .sdmatte import prompts as P
|
||
|
||
model = sdmatte_model
|
||
device = model.device
|
||
dtype = next(model.unet.parameters()).dtype
|
||
size = int(inference_size)
|
||
|
||
B, H, W, C = image.shape
|
||
|
||
# ComfyUI 的 IMAGE 约定是 RGB，但抠图类节点（BiRefNet / RMBG 等）常输出 RGBA。
|
||
# SDMatte 的 VAE 编码器只收 3 通道，多出的 alpha 必须丢掉而不能当颜色喂进去。
|
||
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()
|
||
|
||
# 掩码可能与图像批次数不一致，按 ComfyUI 惯例广播
|
||
if mask.dim() == 2:
|
||
mask = mask.unsqueeze(0)
|
||
if mask.shape[0] != B:
|
||
mask = mask[:1].repeat(B, 1, 1)
|
||
|
||
# 提示类型只是 forward 里的一个分支选择，切换它无需重建模型
|
||
model.aux_input = prompt_type
|
||
|
||
images_t, aux_t, coords_t = [], [], []
|
||
coor_name = None
|
||
|
||
for b in range(B):
|
||
img_np = image[b].detach().cpu().float().numpy() # [H,W,3] in [0,1]
|
||
msk_np = mask[b].detach().cpu().float().numpy() # [H,W] in [0,1]
|
||
|
||
img_r = P.resize_image(img_np, size)
|
||
# 官方 Resize 对 alpha 用双线性；GenMask 的既有掩码分支用最近邻
|
||
interp = cv2.INTER_NEAREST if prompt_type == "mask" else cv2.INTER_LINEAR
|
||
msk_r = cv2.resize(msk_np, (size, size), interpolation=interp)
|
||
msk_r = np.clip(msk_r, 0.0, 1.0)
|
||
|
||
aux_np, coords_np = P.build_prompt(
|
||
msk_r, prompt_type, point_radius=point_radius, seed=seed + b
|
||
)
|
||
|
||
images_t.append(torch.from_numpy(P.normalize(img_r)).permute(2, 0, 1))
|
||
aux_t.append(torch.from_numpy(P.normalize(aux_np)).unsqueeze(0))
|
||
coords_t.append(torch.from_numpy(coords_np))
|
||
|
||
from .sdmatte.meta_arch import AUX_INPUT_DIT
|
||
coor_name = AUX_INPUT_DIT[prompt_type]
|
||
|
||
data = {
|
||
"image": torch.stack(images_t).to(device=device, dtype=dtype),
|
||
prompt_type: torch.stack(aux_t).to(device=device, dtype=dtype),
|
||
coor_name: torch.stack(coords_t).to(device=device, dtype=dtype),
|
||
"is_trans": torch.tensor([1 if is_transparent else 0] * B, dtype=torch.long),
|
||
"caption": [caption] * B,
|
||
}
|
||
|
||
with torch.no_grad():
|
||
pred = model(data) # [B,1,size,size] in [0,1]
|
||
|
||
pred = pred.detach().float().cpu()
|
||
|
||
# 缩放回原图尺寸。官方 inference.py 用 cv2 双线性，并会量化到 uint8；
|
||
# 这里保留浮点，避免白白丢掉 8bit 之外的过渡信息。
|
||
alphas = []
|
||
for b in range(B):
|
||
a = pred[b, 0].numpy()
|
||
a = cv2.resize(a, (W, H), interpolation=cv2.INTER_LINEAR)
|
||
alphas.append(torch.from_numpy(np.clip(a, 0.0, 1.0)))
|
||
alpha = torch.stack(alphas) # [B,H,W]
|
||
|
||
# SDMatte*（即 SDMatte_plus）的训练集以 COCO-Matte 替换了 RefMatte，未学过点提示，
|
||
# 喂 point_mask 会输出接近全黑。这不会报错，只会悄悄给出空结果，故主动提示。
|
||
if prompt_type == "point_mask" and float(alpha.max()) < 0.1:
|
||
print(
|
||
"[Ruinode-SDMatte] 警告：point_mask 提示下输出接近全黑。"
|
||
"官方 SDMatte_plus(SDMatte*) 不支持点提示（其训练集用 COCO-Matte 替换了 RefMatte）。"
|
||
"请改用 bbox_mask，或换成 SDMatte.pth。"
|
||
)
|
||
|
||
cutout = image.detach().cpu().float() * alpha.unsqueeze(-1)
|
||
|
||
return (alpha, cutout)
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"RuiSDMatteLoader": RuiSDMatteLoader,
|
||
"RuiSDMatte": RuiSDMatte,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"RuiSDMatteLoader": "SDMatte 加载器",
|
||
"RuiSDMatte": "SDMatte 精细抠图",
|
||
}
|