rui40000andClaude Opus 4.8
e1a575fcda
基于 SDMatte（vivo 相机研究院，ICCV 2025）的交互式抠图节点， 擅长发丝、绒毛、玻璃、烟雾等常规抠图模型处理不好的边缘。 实现要点： - 严格照搬官方 configs/SDMatte.py 的推理配置（bbox 视觉提示、fp32、1024 分辨率）， 不做任何启发式后处理，输出即模型原始 alpha - 内置官方 LongfeiHuang/SDMatte 的配置文件，无需下载 SD 2.1 权重，也无需联网。 官方 load_weight=False 只用 config 搭骨架，全部权重由 checkpoint 覆盖； 原版 SD 2.1 的 config 缺 bbox_time_embed_dim 等三个专有字段，缺字段时直接报错而非猜测 - 同时支持官方 .pth（12.1GB）与社区 .safetensors（5.19GB）。两者模型权重实测 逐像素完全相同，pth 多出的 6.5GB 是 detectron2 的优化器状态；读 pth 时以受限 Unpickler 只解析 model 段，内存占用与 safetensors 相当 - 自动适配 transformers 5.x 移除 text_model 包装层导致的键名漂移， 避免 text_encoder 的 372 个权重被静默丢弃 - 加载后校验 1316 个张量全部对齐，有任何未覆盖/未使用的权重即中止报错 - 默认开启注意力分片，1024 下显存峰值由约 15.5GB 降至 9.1GB，速度反而略快 实测（官方效果图中的羊驼，对比官方公布 alpha）：MAD=0.0113。 同图同权重下 ComfyUI-SDMatte 为 MAD=0.0884，相差 7.8 倍，主因是其 aux_input="trimap" —— 官方 aux_input_list 只含 point_mask/bbox_mask/mask， trimap 从未作为视觉提示参与训练，且该分支坐标恒为 [0,0,1,1]、定位信息丢失。 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
251 lines
11 KiB
Python
251 lines
11 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
SDMatte 主干网络。
|
||
|
||
移植自 vivoCameraResearch/SDMatte 的 modeling/SDMatte/meta_arch.py。
|
||
相对官方版本只做了两类改动，计算图本身逐行保持一致：
|
||
|
||
1. 官方把 .cuda() 硬编码在 forward 各处，这里改为跟随模型自身所在设备，
|
||
以便在 ComfyUI 里支持 CPU / 多卡 / 显存卸载。
|
||
2. 官方 init_submodule 在 load_weight=False 时用 os.path.join 拼子目录，
|
||
这里同样只读配置、不读权重，但对路径做了存在性校验并给出可读的报错。
|
||
|
||
注意 load_weight 恒为 False 是官方推理配置（configs/SDMatte.py）的原意：
|
||
网络结构只从 config.json 构建，全部权重随后由 checkpoint 覆盖，
|
||
因此 Stable Diffusion 2.1 的权重文件在整个流程中不参与。
|
||
"""
|
||
|
||
import os
|
||
import random
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from diffusers import DDIMScheduler, AutoencoderKL
|
||
from diffusers.models.embeddings import get_timestep_embedding
|
||
from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig
|
||
|
||
from .nn_utils import replace_unet_conv_in, replace_attention_mask_method, add_aux_conv_in
|
||
from .replace import CustomUNet
|
||
|
||
AUX_INPUT_DIT = {
|
||
"auto_mask": "auto_coords",
|
||
"point_mask": "point_coords",
|
||
"bbox_mask": "bbox_coords",
|
||
"mask": "mask_coords",
|
||
"trimap": "trimap_coords",
|
||
}
|
||
|
||
|
||
class SDMatte(nn.Module):
|
||
def __init__(
|
||
self,
|
||
pretrained_model_name_or_path,
|
||
conv_scale=3,
|
||
num_inference_steps=1,
|
||
aux_input="bbox_mask",
|
||
use_aux_input=False,
|
||
use_coor_input=True,
|
||
use_dis_loss=True,
|
||
use_attention_mask=True,
|
||
use_encoder_attention_mask=False,
|
||
add_noise=False,
|
||
attn_mask_aux_input=["point_mask", "bbox_mask", "mask"],
|
||
aux_input_list=["point_mask", "bbox_mask", "mask"],
|
||
use_encoder_hidden_states=True,
|
||
residual_connection=False,
|
||
use_attention_mask_list=[True, True, True],
|
||
use_encoder_hidden_states_list=[True, True, True],
|
||
load_weight=True,
|
||
):
|
||
super().__init__()
|
||
self.init_submodule(pretrained_model_name_or_path, load_weight)
|
||
self.num_inference_steps = num_inference_steps
|
||
self.aux_input = aux_input
|
||
self.use_aux_input = use_aux_input
|
||
self.use_coor_input = use_coor_input
|
||
self.use_dis_loss = use_dis_loss
|
||
self.use_attention_mask = use_attention_mask
|
||
self.use_encoder_attention_mask = use_encoder_attention_mask
|
||
self.add_noise = add_noise
|
||
self.attn_mask_aux_input = attn_mask_aux_input
|
||
self.aux_input_list = aux_input_list
|
||
self.use_encoder_hidden_states = use_encoder_hidden_states
|
||
if use_encoder_hidden_states:
|
||
self.unet = add_aux_conv_in(self.unet)
|
||
if not add_noise:
|
||
conv_scale -= 1
|
||
if not use_aux_input:
|
||
conv_scale -= 1
|
||
if conv_scale > 1:
|
||
self.unet = replace_unet_conv_in(self.unet, conv_scale)
|
||
replace_attention_mask_method(self.unet, residual_connection)
|
||
self.text_encoder.requires_grad_(False)
|
||
self.vae.requires_grad_(False)
|
||
self.unet.use_attention_mask_list = use_attention_mask_list
|
||
self.unet.use_encoder_hidden_states_list = use_encoder_hidden_states_list
|
||
|
||
@property
|
||
def device(self):
|
||
return next(self.unet.parameters()).device
|
||
|
||
def init_submodule(self, pretrained_model_name_or_path, load_weight):
|
||
if load_weight:
|
||
self.text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder="text_encoder")
|
||
self.vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae")
|
||
self.unet = CustomUNet.from_pretrained(
|
||
pretrained_model_name_or_path, subfolder="unet", low_cpu_mem_usage=True, ignore_mismatched_sizes=False
|
||
)
|
||
self.noise_scheduler = DDIMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
|
||
self.tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer")
|
||
return
|
||
|
||
# 仅从 config 构建骨架，不加载任何 SD 权重
|
||
unet_path = os.path.join(pretrained_model_name_or_path, "unet")
|
||
unet_cfg_file = os.path.join(unet_path, "config.json")
|
||
if not os.path.isfile(unet_cfg_file):
|
||
raise FileNotFoundError(f"缺少 SDMatte 的 unet 配置：{unet_cfg_file}")
|
||
|
||
unet_config = CustomUNet.load_config(unet_path)
|
||
# SDMatte 在标准 SD 2.1 的 UNet 配置上追加了三个字段，用于坐标嵌入与不透明度嵌入。
|
||
# 缺任何一个都说明配置来源不对（例如误用了原版 SD 2.1 的 config.json），
|
||
# 此时若放任默认值继续，权重会因形状不符被静默丢弃，必须直接报错。
|
||
for field in ("bbox_time_embed_dim", "point_embeddings_input_dim", "bbox_embeddings_input_dim"):
|
||
if field not in unet_config:
|
||
raise ValueError(
|
||
f"unet 配置缺少 SDMatte 专有字段 '{field}'：{unet_cfg_file}\n"
|
||
"这通常是误用了原版 Stable Diffusion 2.1 的 config.json。"
|
||
"请使用 LongfeiHuang/SDMatte 提供的配置文件。"
|
||
)
|
||
|
||
text_config = CLIPTextConfig.from_pretrained(pretrained_model_name_or_path, subfolder="text_encoder")
|
||
self.text_encoder = CLIPTextModel(text_config)
|
||
|
||
vae_path = os.path.join(pretrained_model_name_or_path, "vae")
|
||
self.vae = AutoencoderKL.from_config(AutoencoderKL.load_config(vae_path))
|
||
|
||
self.unet = CustomUNet.from_config(unet_config, low_cpu_mem_usage=True, ignore_mismatched_sizes=False)
|
||
|
||
scheduler_path = os.path.join(pretrained_model_name_or_path, "scheduler", "scheduler_config.json")
|
||
self.noise_scheduler = DDIMScheduler.from_config(DDIMScheduler.load_config(scheduler_path))
|
||
|
||
self.tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer")
|
||
|
||
def forward(self, data):
|
||
device = self.device
|
||
rgb = data["image"].to(device)
|
||
B = rgb.shape[0]
|
||
|
||
if self.aux_input is None and self.training:
|
||
aux_input_type = random.choice(self.aux_input_list)
|
||
elif self.aux_input is None:
|
||
aux_input_type = "point_mask"
|
||
else:
|
||
aux_input_type = self.aux_input
|
||
|
||
# 视觉提示 -> latent
|
||
if self.use_aux_input:
|
||
aux_input = data[aux_input_type].to(device)
|
||
aux_input = aux_input.repeat(1, 3, 1, 1)
|
||
aux_input_h = self.vae.encoder(aux_input.to(rgb.dtype))
|
||
aux_input_moments = self.vae.quant_conv(aux_input_h)
|
||
aux_input_mean, _ = torch.chunk(aux_input_moments, 2, dim=1)
|
||
aux_input_latent = aux_input_mean * self.vae.config.scaling_factor
|
||
else:
|
||
aux_input_latent = None
|
||
|
||
# 视觉提示的坐标嵌入
|
||
coor_name = AUX_INPUT_DIT[aux_input_type]
|
||
coor = data[coor_name].to(device)
|
||
if coor_name == "point_coords":
|
||
N = coor.shape[1]
|
||
for i in range(N, 1680):
|
||
if 1680 % i == 0:
|
||
num_channels = 1680 // i
|
||
pad_size = i - N
|
||
padding = torch.zeros((B, pad_size), dtype=coor.dtype, device=coor.device)
|
||
coor = torch.cat([coor, padding], dim=1)
|
||
zero_coor = torch.zeros((B, pad_size + N), dtype=coor.dtype, device=coor.device)
|
||
break
|
||
if self.use_coor_input:
|
||
coor = get_timestep_embedding(
|
||
coor.flatten(),
|
||
num_channels,
|
||
flip_sin_to_cos=True,
|
||
downscale_freq_shift=0,
|
||
)
|
||
else:
|
||
coor = get_timestep_embedding(
|
||
zero_coor.flatten(),
|
||
num_channels,
|
||
flip_sin_to_cos=True,
|
||
downscale_freq_shift=0,
|
||
)
|
||
added_cond_kwargs = {"point_coords": coor}
|
||
else:
|
||
if self.use_coor_input:
|
||
added_cond_kwargs = {"bbox_mask_coords": coor}
|
||
else:
|
||
coor = torch.tensor([[0, 0, 1, 1]] * B, device=device)
|
||
added_cond_kwargs = {"bbox_mask_coords": coor}
|
||
|
||
# 掩码自注意力
|
||
if self.use_attention_mask and aux_input_type in self.attn_mask_aux_input:
|
||
attention_mask = data[aux_input_type].to(device)
|
||
attention_mask = (attention_mask + 1) / 2
|
||
attention_mask = F.interpolate(attention_mask, scale_factor=1 / 8, mode="nearest")
|
||
attention_mask = attention_mask.flatten(start_dim=1)
|
||
else:
|
||
attention_mask = None
|
||
|
||
# 原图 -> latent
|
||
rgb_h = self.vae.encoder(rgb)
|
||
rgb_moments = self.vae.quant_conv(rgb_h)
|
||
rgb_mean, _ = torch.chunk(rgb_moments, 2, dim=1)
|
||
rgb_latent = rgb_mean * self.vae.config.scaling_factor
|
||
|
||
# 视觉提示驱动的交互条件
|
||
encoder_hidden_states = None
|
||
if self.use_encoder_hidden_states and aux_input_latent is not None:
|
||
encoder_hidden_states = self.unet.aux_conv_in(aux_input_latent)
|
||
encoder_hidden_states = encoder_hidden_states.view(B, 1024, -1)
|
||
encoder_hidden_states = encoder_hidden_states.permute(0, 2, 1)
|
||
|
||
if "caption" in data:
|
||
prompt = data["caption"]
|
||
else:
|
||
prompt = [""] * B
|
||
prompt = [prompt] if isinstance(prompt, str) else prompt
|
||
text_inputs = self.tokenizer(
|
||
prompt,
|
||
padding="max_length",
|
||
max_length=self.tokenizer.model_max_length,
|
||
truncation=True,
|
||
return_tensors="pt",
|
||
)
|
||
text_input_ids = text_inputs.input_ids.to(device)
|
||
text_embed = self.text_encoder(text_input_ids)[0]
|
||
encoder_hidden_states_2 = text_embed
|
||
|
||
# 不透明度嵌入：透明物体走另一条分支
|
||
is_trans = data["is_trans"].to(device)
|
||
trans = 1 - is_trans
|
||
|
||
# 官方在此处构造了 timestep 张量却传入 None，即单步、不加噪。此处保持一致。
|
||
label_latent = self.unet(
|
||
sample=torch.cat([rgb_latent, aux_input_latent], dim=1),
|
||
trans=trans,
|
||
timestep=None,
|
||
encoder_hidden_states=encoder_hidden_states,
|
||
encoder_hidden_states_2=encoder_hidden_states_2,
|
||
added_cond_kwargs=added_cond_kwargs,
|
||
attention_mask=attention_mask,
|
||
).sample
|
||
label_latent = label_latent / self.vae.config.scaling_factor
|
||
z = self.vae.post_quant_conv(label_latent)
|
||
stacked = self.vae.decoder(z)
|
||
label_mean = stacked.mean(dim=1, keepdim=True)
|
||
output = torch.clip(label_mean, -1.0, 1.0)
|
||
output = (output + 1.0) / 2.0
|
||
return output
|