diff --git a/README.md b/README.md index ec54aef..5788dbf 100644 --- a/README.md +++ b/README.md @@ -5,11 +5,18 @@ ## 版本信息 -- 当前版本：v1.2 +- 当前版本：v0.2 - 更新日期：2024-06-XX ## 更新日志 +### v0.2 (2024-06-XX) +- BiRefNetUltraV3_RBS 节点合并模型加载功能，直接选择模型版本，无需单独加载节点。 +- BiRefNetUltraV3_RBS、TransparentBackgroundUltra_RBS 节点的 max_megapixels 参数支持动态调整分辨率，自动缩放为32的倍数，防止patch分割报错。 +- 新增 Mask Process Details (RBS) 节点，支持VITMatte/GuidedFilter等多种细节处理方法，参数可调。 +- 节点命名与界面优化，部分节点重命名。 +- 细节处理参数范围优化。 + ### v1.2 (2024-06-XX) - 节点界面极简化，移除所有细节处理相关参数和功能。 - TransparentBackgroundUltra_RBS 去除 WIP 标志。 @@ -25,36 +32,41 @@ ## 节点说明 -### 1. LoadBiRefNetModelV3_RBS（推荐） -- **功能**：自动下载并加载 BiRefNet 最新模型，支持 BiRefNet_dynamic。 -- **参数**： - - `version`：选择模型版本（BiRefNet-General、RMBG-2.0、BiRefNet_dynamic）。 -- **输出**：`birefnet_model`（供后续节点使用） - -### 2. BiRefNetUltraV3_RBS（推荐） -- **功能**：使用 BiRefNet Ultra V3 进行高质量背景移除，支持 dynamic 动态模型。 +### 1. BiRefNet Ultra V3 (RBS) +- **功能**：使用 BiRefNet Ultra V3 进行高质量背景移除，支持 dynamic/HR/HR-matting 等多种模型。 - **参数**： - `image`：输入图片（支持批量）。 - - `birefnet_model`：已加载的模型（支持 dynamic）。 - - `device`、`max_megapixels`：详见节点界面。 + - `version`：选择模型版本（BiRefNet-General、RMBG-2.0、BiRefNet_dynamic等）。 + - `device`、`max_megapixels`：详见节点界面。max_megapixels 会自动动态缩放分辨率为32的倍数。 - **输出**： - `image`：去背景后的 RGBA 图片 - `mask`：前景掩码 -### 3. TransparentBackgroundUltra_RBS +### 2. Transparent Background Ultra (RBS) - **功能**：将图片背景转换为透明。 - **参数**： - `image`：输入图片。 - `model`：选择本地模型。 - - `device`、`max_megapixels`：详见节点界面。 + - `device`、`max_megapixels`：详见节点界面。max_megapixels 会自动动态缩放分辨率为32的倍数。 - **输出**： - `image`：透明背景图片 - `mask`：前景掩码 +### 3. Mask Process Details (RBS) +- **功能**：对输入图像和mask进行细节优化，支持VITMatte、GuidedFilter等多种方法。 +- **参数**： + - `image`：输入图片。 + - `mask`：输入掩码。 + - `detail_method`：细节处理方法（VITMatte、VITMatte(local)、GuidedFilter）。 + - `detail_erode`、`detail_dilate`、`black_point`、`white_point`、`device`、`max_megapixels`：详见节点界面。 +- **输出**： + - `image`：优化后的图片 + - `mask`：优化后的掩码 + ## 典型用法 -1. 用 `LoadBiRefNetModelV3_RBS` 加载模型（推荐选择 BiRefNet_dynamic）。 -2. 用 `BiRefNetUltraV3_RBS` 进行背景移除。 -3. 可选：用 `TransparentBackgroundUltra_RBS` 进一步处理透明背景。 +1. 用 `BiRefNet Ultra V3 (RBS)` 进行背景移除。 +2. 可选：用 `Mask Process Details (RBS)` 进一步优化边缘细节。 +3. 可选：用 `Transparent Background Ultra (RBS)` 处理透明背景。 ## 注意事项 - 请将模型文件放在 `ComfyUI/models/transparent-background/` 目录下，或使用新版节点自动下载。 diff --git a/__init__.py b/__init__.py index 825a928..686c7ba 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,5 @@ from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS -__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] + +__version__ = 'v0.2' \ No newline at end of file diff --git a/nodes.py b/nodes.py index d5609d5..360e0e6 100644 --- a/nodes.py +++ b/nodes.py @@ -74,7 +74,26 @@ class TransparentBackgroundUltra_RBS: for i in image: i = torch.unsqueeze(i, 0) orig_image = tensor2pil(i).convert('RGB') - ret_image = remover.process(orig_image, type='rgba') + width, height = orig_image.size + max_pixels = int(max_megapixels * 1_048_576) + orig_pixels = width * height + patch_size = 32 + if orig_pixels > max_pixels: + scale = (max_pixels / orig_pixels) ** 0.5 + new_width = max(1, int(width * scale)) + new_height = max(1, int(height * scale)) + else: + new_width, new_height = width, height + new_width = (new_width // patch_size) * patch_size + new_height = (new_height // patch_size) * patch_size + new_width = max(patch_size, new_width) + new_height = max(patch_size, new_height) + inference_image_size = (new_width, new_height) + log(f"[TransparentBackgroundUltra_RBS] 原始尺寸: {width}x{height}, 实际推理尺寸: {inference_image_size[0]}x{inference_image_size[1]}", message_type='info') + resized_image = orig_image.resize(inference_image_size, Image.BILINEAR) + ret_image = remover.process(resized_image, type='rgba') + # 推理后还原为原图尺寸 + ret_image = ret_image.resize(orig_image.size, Image.BILINEAR) _mask = ret_image.split()[3] ret_images.append(pil2tensor(ret_image)) ret_masks.append(image2mask(_mask)) @@ -83,26 +102,7 @@ class TransparentBackgroundUltra_RBS: # ======================== V3 新增：支持 BiRefNet-Dynamic ======================== -class LoadBiRefNetModelV3_RBS: - def __init__(self): - self.model = None - - @classmethod - def INPUT_TYPES(s): - # 新增 dynamic 选项 - model_list = list(s.birefnet_model_repos.keys()) - return { - "required": { - "version": (model_list, {"default": model_list[0]}), # 选择模型版本 - }, - } - - RETURN_TYPES = ("BIREFNET_MODEL",) - RETURN_NAMES = ("birefnet_model",) - FUNCTION = "load_birefnet_model_v3" - CATEGORY = 'RemoveBackgroundSuite' - - # 支持 dynamic、HR、HR-matting 模型 +class BiRefNetUltraV3_RBS: birefnet_model_repos = { "BiRefNet-General": "ZhengPeng7/BiRefNet", "RMBG-2.0": "briaai/RMBG-2.0", @@ -111,48 +111,18 @@ class LoadBiRefNetModelV3_RBS: "BiRefNet_HR-matting": "ZhengPeng7/BiRefNet_HR-matting" } - def load_birefnet_model_v3(self, version): - birefnet_path = os.path.join(folder_paths.models_dir, 'BiRefNet') - os.makedirs(birefnet_path, exist_ok=True) - model_path = os.path.join(birefnet_path, version) - - if version == "BiRefNet-General": - old_birefnet_path = os.path.join(birefnet_path, 'pth') - old_model = "BiRefNet-general-epoch_244.pth" - old_model_path = os.path.join(old_birefnet_path, old_model) - if os.path.exists(old_model_path): - from .BiRefNet.models.birefnet import BiRefNet - from .BiRefNet.utils import check_state_dict - self.birefnet = BiRefNet(bb_pretrained=False) - self.state_dict = torch.load(old_model_path, map_location='cpu', weights_only=True) - self.state_dict = check_state_dict(self.state_dict) - self.birefnet.load_state_dict(self.state_dict) - return (self.birefnet,) - # 动态模型及HR、HR-matting模型下载 - if version in ["BiRefNet_dynamic", "BiRefNet_HR", "BiRefNet_HR-matting"] and not os.path.exists(model_path): - log(f"Downloading {version} model...") - from huggingface_hub import snapshot_download - repo_id = self.birefnet_model_repos[version] - snapshot_download(repo_id=repo_id, local_dir=model_path, ignore_patterns=["*.md", "*.txt"]) - elif version == "RMBG-2.0" and not os.path.exists(model_path): - log(f"Downloading RMBG-2.0 model...") - from huggingface_hub import snapshot_download - snapshot_download(repo_id="briaai/RMBG-2.0", local_dir=model_path, ignore_patterns=["*.md", "*.txt"]) - - self.model = AutoModelForImageSegmentation.from_pretrained(model_path, trust_remote_code=True) - return (self.model,) - -class BiRefNetUltraV3_RBS: def __init__(self): self.NODE_NAME = 'BiRefNetUltraV3_RBS' + self.model_cache = {} @classmethod def INPUT_TYPES(cls): device_list = ['cuda', 'cpu'] + model_list = list(cls.birefnet_model_repos.keys()) return { "required": { "image": ("IMAGE",), - "birefnet_model": ("BIREFNET_MODEL",), + "version": (model_list, {"default": model_list[0]}), "device": (device_list,), "max_megapixels": ("FLOAT", {"default": 2.0, "min": 1, "max": 999, "step": 0.1}), }, @@ -165,12 +135,45 @@ class BiRefNetUltraV3_RBS: FUNCTION = "birefnet_ultra_v3" CATEGORY = 'RemoveBackgroundSuite' - # 主推理流程，兼容 dynamic 模型 - def birefnet_ultra_v3(self, image, birefnet_model, device, max_megapixels): + def load_birefnet_model(self, version): + birefnet_path = os.path.join(folder_paths.models_dir, 'BiRefNet') + os.makedirs(birefnet_path, exist_ok=True) + model_path = os.path.join(birefnet_path, version) + + if version == "BiRefNet-General": + old_birefnet_path = os.path.join(birefnet_path, 'pth') + old_model = "BiRefNet-general-epoch_244.pth" + old_model_path = os.path.join(old_birefnet_path, old_model) + if os.path.exists(old_model_path): + from .BiRefNet.models.birefnet import BiRefNet + from .BiRefNet.utils import check_state_dict + birefnet = BiRefNet(bb_pretrained=False) + state_dict = torch.load(old_model_path, map_location='cpu', weights_only=True) + state_dict = check_state_dict(state_dict) + birefnet.load_state_dict(state_dict) + return birefnet + # 动态模型及HR、HR-matting模型下载 + if version in ["BiRefNet_dynamic", "BiRefNet_HR", "BiRefNet_HR-matting"] and not os.path.exists(model_path): + log(f"Downloading {version} model...") + from huggingface_hub import snapshot_download + repo_id = self.birefnet_model_repos[version] + snapshot_download(repo_id=repo_id, local_dir=model_path, ignore_patterns=["*.md", "*.txt"]) + elif version == "RMBG-2.0" and not os.path.exists(model_path): + log(f"Downloading RMBG-2.0 model...") + from huggingface_hub import snapshot_download + snapshot_download(repo_id="briaai/RMBG-2.0", local_dir=model_path, ignore_patterns=["*.md", "*.txt"]) + + model = AutoModelForImageSegmentation.from_pretrained(model_path, trust_remote_code=True) + return model + + def birefnet_ultra_v3(self, image, version, device, max_megapixels): ret_images = [] ret_masks = [] - inference_image_size = (1024, 1024) torch.set_float32_matmul_precision(['high', 'highest'][0]) + # 模型缓存，避免重复加载 + if version not in self.model_cache: + self.model_cache[version] = self.load_birefnet_model(version) + birefnet_model = self.model_cache[version] birefnet_model.to(device) birefnet_model.eval() comfy_pbar = ProgressBar(len(image)) @@ -178,6 +181,24 @@ class BiRefNetUltraV3_RBS: for i in image: i = torch.unsqueeze(i, 0) orig_image = tensor2pil(i).convert('RGB') + width, height = orig_image.size + max_pixels = int(max_megapixels * 1_048_576) + orig_pixels = width * height + patch_size = 32 # 保证分辨率为32的倍数 + # 动态调整分辨率 + if orig_pixels > max_pixels: + scale = (max_pixels / orig_pixels) ** 0.5 + new_width = max(1, int(width * scale)) + new_height = max(1, int(height * scale)) + else: + new_width, new_height = width, height + # 向下取整为patch_size的倍数 + new_width = (new_width // patch_size) * patch_size + new_height = (new_height // patch_size) * patch_size + new_width = max(patch_size, new_width) + new_height = max(patch_size, new_height) + inference_image_size = (new_width, new_height) + log(f"[BiRefNetUltraV3_RBS] 原始尺寸: {width}x{height}, 实际推理尺寸: {inference_image_size[0]}x{inference_image_size[1]}", message_type='info') transform_image = transforms.Compose([ transforms.Resize(inference_image_size), transforms.ToTensor(), @@ -188,9 +209,9 @@ class BiRefNetUltraV3_RBS: preds = birefnet_model(inference_image)[-1].sigmoid().cpu() pred = preds[0].squeeze() pred_pil = transforms.ToPILImage()(pred) - _mask = pred_pil.resize(inference_image_size) - resize_sampler = Image.BILINEAR - _mask = _mask.resize(orig_image.size, resize_sampler) + # 先还原到inference_image_size，再还原到原图尺寸 + _mask = pred_pil.resize(inference_image_size, Image.BILINEAR) + _mask = _mask.resize(orig_image.size, Image.BILINEAR) brightness_image = ImageEnhance.Brightness(_mask) _mask = brightness_image.enhance(factor=1.08) _mask = image2mask(_mask) @@ -208,17 +229,68 @@ class BiRefNetUltraV3_RBS: log(f"{self.NODE_NAME} Processed {len(ret_masks)} image(s).", message_type='finish') return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),) +class ProcessDetails_RBS: + def __init__(self): + self.NODE_NAME = 'ProcessDetails_RBS' + self.vitmatte_model = None + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "mask": ("MASK",), + "detail_method": (["VITMatte", "VITMatte(local)", "GuidedFilter"], {"default": "VITMatte"}), + "detail_erode": ("INT", {"default": 4, "min": 1, "max": 100, "step": 1}), + "detail_dilate": ("INT", {"default": 2, "min": 1, "max": 100, "step": 1}), + "black_point": ("FLOAT", {"default": 0.01, "min": 0.01, "max": 0.98, "step": 0.01}), + "white_point": ("FLOAT", {"default": 0.99, "min": 0.02, "max": 0.99, "step": 0.01}), + "device": (["cuda", "cpu"], {"default": "cuda"}), + "max_megapixels": ("FLOAT", {"default": 2.0, "min": 1, "max": 999, "step": 0.1}), + }, + } + + RETURN_TYPES = ("IMAGE", "MASK",) + RETURN_NAMES = ("image", "mask",) + FUNCTION = "process_details" + CATEGORY = 'RemoveBackgroundSuite' + + def process_details(self, image, mask, detail_method, detail_erode, detail_dilate, black_point, white_point, device, max_megapixels): + ret_images = [] + ret_masks = [] + + for i, m in zip(image, mask): + orig_image = tensor2pil(i) + orig_mask = tensor2pil(m) + + if detail_method.startswith("VITMatte"): + local_files_only = detail_method == "VITMatte(local)" + trimap = generate_VITMatte_trimap(orig_mask, detail_erode, detail_dilate) + processed_mask = generate_VITMatte(orig_image, trimap, local_files_only, device, max_megapixels) + else: # GuidedFilter + processed_mask = mask_edge_detail(i, m, detail_erode, black_point, white_point) + processed_mask = tensor2pil(processed_mask) + + # 应用处理后的蒙版到图像 + processed_image = RGB2RGBA(orig_image, processed_mask) + + ret_images.append(pil2tensor(processed_image)) + ret_masks.append(image2mask(processed_mask)) + + log(f"{self.NODE_NAME} Processed {len(ret_images)} image(s).", message_type='finish') + return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),) + # 节点注册映射 NODE_CLASS_MAPPINGS = { "TransparentBackgroundUltra_RBS": TransparentBackgroundUltra_RBS, - "LoadBiRefNetModelV3_RBS": LoadBiRefNetModelV3_RBS, - "BiRefNetUltraV3_RBS": BiRefNetUltraV3_RBS + "BiRefNetUltraV3_RBS": BiRefNetUltraV3_RBS, + "MaskProcessDetails_RBS": ProcessDetails_RBS } NODE_DISPLAY_NAME_MAPPINGS = { "TransparentBackgroundUltra_RBS": "Transparent Background Ultra (RBS)", - "LoadBiRefNetModelV3_RBS": "Load BiRefNet Model V3 (RBS)", - "BiRefNetUltraV3_RBS": "BiRefNet Ultra V3 (RBS)" + "BiRefNetUltraV3_RBS": "BiRefNet Ultra V3 (RBS)", + "MaskProcessDetails_RBS": "Mask Process Details (RBS)" } def log_mask_info(tag, mask):