Optimize overseas RMBG responses to return validated masks only

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 <KEY>`。未配置时不附加鉴权头。
 
+## 海外 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.
 

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],
         }

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))

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