From f40be8ec1fbabe1e3cf99d2e8c4eae07a3418c98 Mon Sep 17 00:00:00 2001 From: xuyucai Date: Tue, 3 Nov 2026 15:24:02 +0800 Subject: [PATCH] Optimize overseas RMBG responses to return validated masks only --- README.md | 8 ++ .../zhiyi_rmbg_segment_node.py | 47 +++++- tests/test_rmbg_mask_only.py | 135 ++++++++++++++++++ tests/test_rmbg_segment_node.py | 7 +- 4 files changed, 193 insertions(+), 4 deletions(-) create mode 100644 tests/test_rmbg_mask_only.py diff --git a/README.md b/README.md index cfec312..36ed79a 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,14 @@ URL 模式沿用现有 OSS 上传配置：`FD_OSS_BUCKET_NAME`、`FD_OSS_ENDPOIN 若模型服务开启鉴权，分别配置 `FD_SAM2_API_KEY` 和 `FD_DWPOSE_API_KEY`，节点会对推理与健康检查请求附加 `Authorization: Bearer `。未配置时不附加鉴权头。 +## 海外 RMBG 仅返回 mask + +`ZhiYiRMBGNode` / `ZhiYiRMBGNodeV2` 调用已配置的海外 RMBG 网关时，默认只请求 `return_image=false`、`return_mask=true`、`return_mask_image=false`。国内服务沿用无损 8 位灰度 PNG + Base64 返回 mask；海外用节点保留的原始 RGB 和 mask 合成 Alpha 透明图或 Color 背景换色图，MASK_IMAGE 也在本地生成。INFO 中的三个 return 字段记录实际请求值，`mask_only_response` 标明此路径是否启用。 + +此优化只匹配 `FD_OVERSEAS_GPU_GATEWAY` 下的 `model-api-dwpose/v1/rmbg`，覆盖旧版 service_url 与 V2 海外环境，其他分割服务和国内地址沿用原返回参数。无需修改工作流；回退时设置 `FD_OVERSEAS_RMBG_MASK_ONLY=0` 并重启 ComfyUI。该开关独立于图片输入编码，本次继续传原尺寸 PNG。 + +仅 mask 模式要求返回 `image_size` 与原图一致，灰度 PNG 的实际尺寸与 `output_size` 一致，且输出尺寸属于原图尺寸或声明的 `inference_size`。推理尺寸的 mask 按现有 BILINEAR 规则恢复原尺寸；缺失、损坏、错误尺寸、非灰度 PNG 或服务仍返回成图均明确报错，不用空 mask 或成图兜底。该方式保持当前节点使用原图合成的语义，不新增前景颜色细化。 + > [!NOTE] > This projected was created with a [cookiecutter](https://github.com/Comfy-Org/cookiecutter-comfy-extension) template. It helps you start writing custom nodes without worrying about the Python setup. diff --git a/src/Comfyui_Fd_Nodes/zhiyi_rmbg_segment_node.py b/src/Comfyui_Fd_Nodes/zhiyi_rmbg_segment_node.py index a9cdd30..d49800a 100644 --- a/src/Comfyui_Fd_Nodes/zhiyi_rmbg_segment_node.py +++ b/src/Comfyui_Fd_Nodes/zhiyi_rmbg_segment_node.py @@ -2,6 +2,7 @@ import base64 import io import json import logging +import os import time import traceback from concurrent.futures import ThreadPoolExecutor, as_completed @@ -18,6 +19,7 @@ from .config import ( FD_CLOTHES_SEGMENT_URL, FD_DWPOSE_API_KEY, FD_FASHION_SEGMENT_URL, + FD_OVERSEAS_DWPOSE_URL, FD_RMBG_URL, RMBG_SERVICE_URL_PRESETS, ) @@ -41,6 +43,17 @@ DEFAULT_CLOTHES_SEGMENT_URL = "http://10.1.0.230:8003/v1/segment/clothes" DEFAULT_FASHION_SEGMENT_URL = "http://10.1.0.230:8003/v1/segment/fashion" DEFAULT_BODY_SEGMENT_URL = "http://10.1.0.230:8003/v1/segment/body" BACKGROUND_MODES = ["Alpha", "Color"] +FD_OVERSEAS_RMBG_MASK_ONLY = os.getenv("FD_OVERSEAS_RMBG_MASK_ONLY", "1") == "1" + + +def _response_size(result: dict, field: str) -> tuple[int, int]: + value = result.get(field) + if not isinstance(value, dict): + raise RuntimeError(f"RMBG mask 响应缺少 {field} 尺寸") + size = (value.get("width"), value.get("height")) + if any(type(dimension) is not int or dimension <= 0 for dimension in size): + raise RuntimeError(f"RMBG mask 响应 {field} 尺寸无效: {value}") + return size CLOTHES_CLASSES = [ "Background", @@ -243,6 +256,27 @@ class _RmbgSegmentApiBase: mask_array = np.array(mask_image).astype(np.float32) / 255.0 return torch.from_numpy(np.clip(mask_array, 0.0, 1.0).astype(np.float32)).unsqueeze(0) + def _rmbg_mask_only_to_tensor(self, result: dict, source_size: tuple[int, int]) -> torch.Tensor: + """Validate the native mask before restoring the service's inference size.""" + mask_image = self._decode_png_data_url(result.get("mask"), "mask") + if mask_image.format != "PNG" or mask_image.mode != "L": + raise RuntimeError("RMBG mask 必须为 8 位灰度 PNG") + image_size = _response_size(result, "image_size") + inference_size = _response_size(result, "inference_size") + output_size = _response_size(result, "output_size") + if image_size != source_size: + raise RuntimeError(f"RMBG image_size 与原图尺寸不匹配: {image_size} != {source_size}") + if mask_image.size != output_size or output_size not in (source_size, inference_size): + raise RuntimeError( + f"RMBG mask 尺寸异常: mask={mask_image.size}, output_size={output_size}, " + f"image_size={image_size}, inference_size={inference_size}" + ) + if result.get("processed_image") or result.get("mask_image"): + raise RuntimeError("RMBG 服务未遵守仅返回 mask 的请求，仍返回 processed_image 或 mask_image") + if mask_image.size != source_size: + mask_image = mask_image.resize(source_size, Image.Resampling.BILINEAR) + return torch.from_numpy(np.array(mask_image).astype(np.float32) / 255.0).unsqueeze(0) + def _empty_mask_tensor(self, size): width, height = size return torch.zeros((1, height, width), dtype=torch.float32) @@ -485,7 +519,9 @@ class _RmbgSegmentApiBase: logger.info("Calling %s API idx=%s service_url=%s", label, idx, service_url) result = self._post_process(service_url, payload, timeout, label) - if result.get("mask"): + if payload_params["mask_only_response"]: + mask = self._rmbg_mask_only_to_tensor(result, original_image.size) + elif result.get("mask"): mask = self._mask_data_url_to_tensor(result["mask"], original_image.size) else: mask = self._empty_mask_tensor(original_image.size) @@ -578,6 +614,13 @@ class _RmbgSegmentApiBase: RMBG_SERVICE_URL_PRESETS, ) final_service_url = _normalize_rmbg_segment_url(selected_service_url, default_url, endpoint) + mask_only_response = ( + FD_OVERSEAS_RMBG_MASK_ONLY + and endpoint == "rmbg" + and final_service_url == f"{FD_OVERSEAS_DWPOSE_URL}/v1/rmbg" + ) + if mask_only_response: + return_image, return_mask, return_mask_image = False, True, False image_tensors = self._expand_images(image) if not image_tensors: raise RuntimeError("未提供图片") @@ -595,6 +638,7 @@ class _RmbgSegmentApiBase: "return_image": return_image, "return_mask": return_mask, "return_mask_image": return_mask_image, + "mask_only_response": mask_only_response, "extra": extra, } tasks = [] @@ -619,6 +663,7 @@ class _RmbgSegmentApiBase: "return_image": bool(return_image), "return_mask": bool(return_mask), "return_mask_image": bool(return_mask_image), + "mask_only_response": mask_only_response, **info_extra, "items": [result["info"] for result in results], } diff --git a/tests/test_rmbg_mask_only.py b/tests/test_rmbg_mask_only.py new file mode 100644 index 0000000..22f8012 --- /dev/null +++ b/tests/test_rmbg_mask_only.py @@ -0,0 +1,135 @@ +import base64 +import io +import json + +import numpy as np +import pytest +import torch +from PIL import Image + +from src.Comfyui_Fd_Nodes import zhiyi_rmbg_segment_node as module +from src.Comfyui_Fd_Nodes.v2.rmbg_segment_v2_node import ZhiYiRMBGNodeV2 + + +def mask_response(size: tuple[int, int] = (5, 3), source_size: tuple[int, int] = (5, 3)) -> dict: + values = np.linspace(0, 255, size[0] * size[1], dtype=np.uint8).reshape(size[1], size[0]) + buffer = io.BytesIO() + Image.fromarray(values).save(buffer, format="PNG") + return { + "mask": "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode("ascii"), + "processed_image": None, + "mask_image": None, + "image_size": {"width": source_size[0], "height": source_size[1]}, + "inference_size": {"width": size[0], "height": size[1]}, + "output_size": {"width": size[0], "height": size[1]}, + } + + +def call_node(node: module._RmbgSegmentApiBase, image: torch.Tensor, **overrides: object) -> tuple: + kwargs = { + name: spec[1]["default"] + for name, spec in node.INPUT_TYPES()["required"].items() + if name != "image" + } + kwargs.update(image=image, service_url=module.FD_OVERSEAS_DWPOSE_URL) + if "env" in kwargs: + kwargs["env"] = "overseas" + kwargs.update(overrides) + return getattr(node, node.FUNCTION)(**kwargs) + + +@pytest.mark.parametrize("node_cls", [module.ZhiYiRMBGNode, ZhiYiRMBGNodeV2]) +@pytest.mark.parametrize("background", ["Alpha", "Color"]) +def test_overseas_returns_only_mask_and_composes_original(monkeypatch, node_cls: type, background: str) -> None: + node = node_cls() + image = torch.arange(45, dtype=torch.float32).reshape(1, 3, 5, 3) / 255 + original = image.clone() + result = mask_response() + captured = [] + + def post(url: str, payload: dict, timeout: int, label: str) -> dict: + captured.append(payload) + assert (payload["return_image"], payload["return_mask"], payload["return_mask_image"]) == (False, True, False) + assert payload["image_format"] == "png_base64" + assert payload["process_res"] == 1024 + return result + + monkeypatch.setattr(node, "_post_process", post) + monkeypatch.setattr(module, "tensor_to_image_payload", lambda *args: "https://example.com/original.png") + output, mask, preview, info_text = call_node( + node, image, background=background, background_color="#123456", + return_image=True, return_mask=False, return_mask_image=True, + ) + alpha = torch.from_numpy(np.linspace(0, 255, 15, dtype=np.uint8).reshape(1, 3, 5).astype(np.float32) / 255) + assert len(captured) == 1 + assert torch.equal(image, original) + assert torch.equal(mask, alpha) + assert torch.equal(preview, alpha.unsqueeze(-1).expand(-1, -1, -1, 3)) + if background == "Alpha": + assert output.shape == (1, 3, 5, 4) + assert torch.equal(output[..., 3], alpha) + assert torch.equal(output[..., :3], image * (alpha.unsqueeze(-1) > 0)) + # Soft edge RGB stays straight alpha; it must not be multiplied twice. + assert torch.equal(output[0, 0, 1, :3], image[0, 0, 1]) + else: + color = torch.tensor([18, 52, 86], dtype=torch.float32) / 255 + assert torch.equal(output, image * alpha.unsqueeze(-1) + color * (1 - alpha.unsqueeze(-1))) + info = json.loads(info_text) + assert info["mask_only_response"] is True + assert (info["return_image"], info["return_mask"], info["return_mask_image"]) == (False, True, False) + + +@pytest.mark.parametrize("disabled,domestic,segment", [(True, False, False), (False, True, False), (False, False, True)]) +def test_scope_and_rollback_keep_original_return_flags(monkeypatch, disabled: bool, domestic: bool, segment: bool) -> None: + node = module.ZhiYiBodySegmentNode() if segment else module.ZhiYiRMBGNode() + monkeypatch.setattr(module, "FD_OVERSEAS_RMBG_MASK_ONLY", not disabled) + monkeypatch.setattr(module, "tensor_to_image_payload", lambda *args: "data:image/png;base64,input") + + def post(url: str, payload: dict, timeout: int, label: str) -> dict: + assert (payload["return_image"], payload["return_mask"], payload["return_mask_image"]) == (True, False, True) + return mask_response() + + monkeypatch.setattr(node, "_post_process", post) + overrides = {"service_url": "http://domestic.example/v1/rmbg"} if domestic else {} + *_, info = call_node(node, torch.zeros(1, 3, 5, 3), return_image=True, return_mask=False, return_mask_image=True, **overrides) + assert json.loads(info)["mask_only_response"] is False + + +def test_declared_inference_mask_is_restored_with_bilinear() -> None: + node = module.ZhiYiRMBGNode() + response = mask_response(size=(2, 2), source_size=(5, 3)) + result = node._rmbg_mask_only_to_tensor(response, (5, 3)) + expected = np.array(Image.fromarray(np.array([[0, 85], [170, 255]], dtype=np.uint8)).resize((5, 3), Image.Resampling.BILINEAR)) + assert torch.equal(result, torch.from_numpy(expected.astype(np.float32) / 255).unsqueeze(0)) + + +@pytest.mark.parametrize("field,value", [ + ("mask", None), ("mask", "broken-base64"), + ("image_size", None), ("image_size", {"width": 3, "height": 5}), + ("image_size", {"width": True, "height": 3}), + ("inference_size", {"width": 0, "height": 3}), + ("output_size", None), ("output_size", {"width": 99, "height": 3}), + ("processed_image", "unexpected-image"), ("mask_image", "unexpected-preview"), +]) +def test_invalid_mask_only_response_fails(field: str, value: object) -> None: + response = mask_response() + response[field] = value + with pytest.raises(RuntimeError, match="mask|image_size|inference_size|output_size"): + module.ZhiYiRMBGNode()._rmbg_mask_only_to_tensor(response, (5, 3)) + + +@pytest.mark.parametrize("mode,format", [("RGB", "PNG"), ("L", "JPEG"), ("I;16", "PNG")]) +def test_mask_encoding_must_be_lossless_grayscale(mode: str, format: str) -> None: + response = mask_response() + buffer = io.BytesIO() + Image.new(mode, (5, 3)).save(buffer, format=format) + response["mask"] = base64.b64encode(buffer.getvalue()).decode("ascii") + with pytest.raises(RuntimeError, match="8 位灰度 PNG"): + module.ZhiYiRMBGNode()._rmbg_mask_only_to_tensor(response, (5, 3)) + + +def test_undeclared_output_size_fails() -> None: + response = mask_response(size=(2, 2), source_size=(5, 3)) + response["inference_size"] = {"width": 4, "height": 2} + with pytest.raises(RuntimeError, match="mask 尺寸异常"): + module.ZhiYiRMBGNode()._rmbg_mask_only_to_tensor(response, (5, 3)) diff --git a/tests/test_rmbg_segment_node.py b/tests/test_rmbg_segment_node.py index 5f52481..68b9c39 100644 --- a/tests/test_rmbg_segment_node.py +++ b/tests/test_rmbg_segment_node.py @@ -319,7 +319,8 @@ def test_rmbg_posts_full_payload_and_decodes_outputs(monkeypatch): assert result_image.shape == (1, 2, 2, 3) assert mask.shape == (1, 2, 2) assert mask_image.shape == (1, 2, 2, 3) - assert torch.equal(result_image[0, 0, 0], torch.tensor([0.0, 1.0, 0.0])) + # The image output uses the original RGB and mask, not the server's green image. + assert torch.equal(result_image[0, 0, 0], torch.tensor([0.0, 0.0, 0.0])) assert torch.equal(mask, torch.tensor([[[1.0, 0.0], [0.0, 1.0]]])) assert torch.equal(mask_image[0, 0, 0], torch.tensor([1.0, 0.0, 0.0])) assert info["service_url"] == "http://example.com:8003/v1/rmbg" @@ -422,7 +423,7 @@ def test_body_health_check_and_process_res_ignored_notice(monkeypatch): assert info["items"][0]["preprocessor"] == "BodySegmentPreprocessor" -def test_node_switch_returns_original_and_empty_mask_without_request(monkeypatch): +def test_node_switch_returns_transparent_image_and_empty_mask_without_request(monkeypatch): node = ZhiYiFashionSegmentNode() image = torch.ones((1, 2, 2, 3), dtype=torch.float32) @@ -450,7 +451,7 @@ def test_node_switch_returns_original_and_empty_mask_without_request(monkeypatch node_switch=1, ) - assert torch.equal(result_image, image) + assert torch.equal(result_image, torch.zeros((1, 2, 2, 4), dtype=torch.float32)) assert torch.equal(mask, torch.zeros((1, 2, 2), dtype=torch.float32)) assert torch.equal(mask_image, torch.zeros((1, 2, 2, 3), dtype=torch.float32)) assert json.loads(info_text)["skipped"] is True