From 52638dfe8c8dfde49f6594eda83840e8b99bc70d Mon Sep 17 00:00:00 2001 From: xuyucai Date: Wed, 4 Nov 2026 20:00:42 +0800 Subject: [PATCH] Use deployment routing and preserve refined foreground in remote RMBG --- README.md | 21 ++-- pyproject.toml | 2 +- src/Comfyui_Fd_Nodes/config.py | 10 ++ .../other/rmbg_workflow_migration.js | 24 +++++ .../v2/rmbg_segment_v2_node.py | 101 +----------------- .../zhiyi_rmbg_segment_node.py | 95 ++++++++-------- tests/test_rmbg_mask_only.py | 1 + tests/test_rmbg_runtime.py | 82 ++++++++++++++ tests/test_rmbg_segment_node.py | 24 +++-- tests/test_v2_nodes.py | 54 +++++----- 10 files changed, 222 insertions(+), 192 deletions(-) create mode 100644 src/Comfyui_Fd_Nodes/other/rmbg_workflow_migration.js create mode 100644 tests/test_rmbg_runtime.py diff --git a/README.md b/README.md index 058a876..73e0644 100644 --- a/README.md +++ b/README.md @@ -21,20 +21,25 @@ URL 模式沿用现有 OSS 上传配置：`FD_OSS_BUCKET_NAME`、`FD_OSS_ENDPOIN 若模型服务开启鉴权，分别配置 `FD_SAM2_API_KEY` 和 `FD_DWPOSE_API_KEY`，节点会对推理与健康检查请求附加 `Authorization: Bearer `。未配置时不附加鉴权头。 -## 远程 RMBG 传输质量（单选） +## 远程 RMBG 部署环境与传输 -`ZhiYiRMBGNode` / `ZhiYiRMBGNodeV2` 在末尾提供可选单选项 `image_encoding`： +`ZhiYiRMBGNode` / `ZhiYiRMBGNodeV2` 使用 Comfy 进程环境变量路由，不再提供 `env`、`service_url`、`service_url_preset` 或三个 `return_*` 节点输入。旧 API JSON 中残留的这些字段会被忽略；不会覆盖部署配置。 -- `png`（默认）：无损 PNG 推理输入；旧工作流缺少该字段时仍使用 PNG。 -- `jpeg_q95`：原尺寸 RGB JPEG 推理副本，quality=95、4:2:0、optimize=true、非 progressive。原输入无论 PNG 或 JPEG 都统一重编码为 Q95。 +| `FD_COMFY_RUNTIME_ENV` | 默认服务 | `image_encoding` 缺省值 | +|---|---|---| +| `dev`（未设置时） | `http://10.1.0.230:8003/v1/rmbg` | `png` | +| `domestic` | `http://model-api-dwpose-svc.online-server-gray:8001/v1/rmbg` | `png` | +| `overseas` | `FD_OVERSEAS_GPU_GATEWAY` 下的 `/model-api-dwpose/v1/rmbg` | `jpeg_q95` | -质量选择与 `env` / `service_url` / `image_delivery` 独立。两种质量均仅请求无损 8 位灰度 PNG mask（`return_image=false`、`return_mask=true`、`return_mask_image=false`），客户端用压缩前原始 RGB 与 mask 合成 Alpha 透明图或 Color 背景换色图，并本地生成 MASK_IMAGE。保留旧 return 入参以兼容 JSON，RMBG 请求固定使用上述实际返回开关；衣物/时尚/身体分割仍沿用原返回参数，不新增质量选项。 +部署可用 `FD_RMBG_URL` 覆盖服务地址（包括生产依赖地址）；非法运行环境在加载配置时明确失败。修改部署变量需重启进程。国内集群与 102 保持无损 PNG，海外默认 JPEG Q95；节点末尾仍可用 `image_encoding=png/jpeg_q95` 手动指定，已保存的显式选择优先于缺省值。Q95 使用原尺寸 RGB、quality=95、4:2:0、optimize=true、非 progressive；有损压缩可能改变 mask 边界和背景残留。 -不再读取 `FD_RMBG_REMOTE_OPTIMIZATION` 或 `FD_OVERSEAS_RMBG_MASK_ONLY`。刷新界面后即可在节点选择质量；无需通过环境变量切换或重启进程。INFO 和正式调用日志记录 `image_encoding`、`mask_only_response`，可直接确认 Q95 是否生效。 +- `refine_foreground=false`：只请求无损 8 位灰度 PNG mask，用压缩前原始 RGB 本地合成 Alpha/Color 输出。 +- `refine_foreground=true`：请求精修图片和 mask，使用服务端精修后的像素；Alpha 保留 RGBA，不用原图覆盖精修结果。服务端漏返图片时明确失败。 +- MASK_IMAGE 始终本地从 mask 生成，不传输重复预览。 -仅 mask 响应要求 `image_size` 与原图一致，灰度 PNG 的实际尺寸与 `output_size` 一致，且输出尺寸属于原图尺寸或声明的 `inference_size`。推理尺寸的 mask 按现有 BILINEAR 规则恢复原尺寸；缺失、损坏、错误尺寸、非灰度 PNG 或服务仍返回成图均明确报错。 +仅 mask 响应校验声明尺寸、原图尺寸、8 位灰度 PNG 编码与返回字段；推理尺寸 mask 按 BILINEAR 恢复原尺寸，异常明确失败。INFO 和调用日志记录编码、返回模式及实际服务地址，INFO 额外记录运行环境。 -选择 Q95 时，`image_delivery=url` 上传内容 MD5 命名的 `.jpg`，base64 模式使用 `data:image/jpeg;base64,`。服务端 `image_format=png_base64` 是响应格式，保持不变。JPEG 有损压缩可能改变分割边界和背景残留，可选择 PNG 比较效果。该选项不自动替换工作流中的本地 RMBG 或 ClothesSegment，款式上身 #313 保持现状。 +保留 `image_delivery=base64/url` 及既有海外 URL 传输约束；Q95 URL 上传 `.jpg`，Base64 使用 `data:image/jpeg;base64,`，服务端响应仍为 PNG。不再读取旧 `FD_RMBG_REMOTE_OPTIMIZATION` / `FD_OVERSEAS_RMBG_MASK_ONLY`。衣物/时尚/身体分割及本地 RMBG 不受本次契约调整影响。 > [!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/pyproject.toml b/pyproject.toml index 3943cf7..59f7c00 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "Comfyui_Fd_Nodes" -version = "53.8.73" # 999开头的版本是debug版 +version = "53.8.74" # 999开头的版本是debug版 description = "Custom nodes for fd project" authors = [ {name = "xuyucai", email = "986327386@qq.com"} diff --git a/src/Comfyui_Fd_Nodes/config.py b/src/Comfyui_Fd_Nodes/config.py index e2dfc83..6358743 100644 --- a/src/Comfyui_Fd_Nodes/config.py +++ b/src/Comfyui_Fd_Nodes/config.py @@ -139,6 +139,16 @@ def _derive_dwpose_preprocess_url(endpoint: str) -> str: FD_CONTROLNET_AUX_LINEART_URL = os.getenv("FD_CONTROLNET_AUX_LINEART_URL") or _derive_dwpose_preprocess_url("lineart") FD_CONTROLNET_AUX_DEPTH_ANYTHING_V2_URL = os.getenv("FD_CONTROLNET_AUX_DEPTH_ANYTHING_V2_URL") or _derive_dwpose_preprocess_url("depth-anything-v2") FD_RMBG_URL = os.getenv("FD_RMBG_URL", "http://10.1.0.230:8003/v1/rmbg") +# RMBG routing belongs to the Comfy process, never to a saved workflow. +FD_COMFY_RUNTIME_ENV = os.getenv("FD_COMFY_RUNTIME_ENV", "dev") +if FD_COMFY_RUNTIME_ENV not in ("dev", "domestic", "overseas"): + raise ValueError("FD_COMFY_RUNTIME_ENV must be dev, domestic or overseas") +RMBG_RUNTIME_URL = os.getenv("FD_RMBG_URL") or { + "dev": "http://10.1.0.230:8003/v1/rmbg", + "domestic": "http://model-api-dwpose-svc.online-server-gray:8001/v1/rmbg", + "overseas": f"{FD_OVERSEAS_DWPOSE_URL}/v1/rmbg", +}[FD_COMFY_RUNTIME_ENV] +RMBG_DEFAULT_IMAGE_ENCODING = "jpeg_q95" if FD_COMFY_RUNTIME_ENV == "overseas" else "png" FD_CLOTHES_SEGMENT_URL = os.getenv("FD_CLOTHES_SEGMENT_URL", "http://10.1.0.230:8003/v1/segment/clothes") FD_FASHION_SEGMENT_URL = os.getenv("FD_FASHION_SEGMENT_URL", "http://10.1.0.230:8003/v1/segment/fashion") FD_BODY_SEGMENT_URL = os.getenv("FD_BODY_SEGMENT_URL", "http://10.1.0.230:8003/v1/segment/body") diff --git a/src/Comfyui_Fd_Nodes/other/rmbg_workflow_migration.js b/src/Comfyui_Fd_Nodes/other/rmbg_workflow_migration.js new file mode 100644 index 0000000..728a2a2 --- /dev/null +++ b/src/Comfyui_Fd_Nodes/other/rmbg_workflow_migration.js @@ -0,0 +1,24 @@ +import { app } from "../../scripts/app.js"; + +// LiteGraph stores widgets positionally. Remove retired fields before loading +// an old workflow so process_res/refine/background/encoding cannot shift. +app.registerExtension({ + name: "ZhiYiRMBG.DeploymentInputs", + async beforeRegisterNodeDef(nodeType, nodeData) { + if (!["ZhiYiRMBGNode", "ZhiYiRMBGNodeV2"].includes(nodeData.name)) return; + const configure = nodeType.prototype.configure; + nodeType.prototype.configure = function (info) { + const values = info.widgets_values; + if (Array.isArray(values) && typeof values[0] === "string") { + const v2 = nodeData.name === "ZhiYiRMBGNodeV2"; + const indices = v2 + ? [2, 3, 4, 5, 6, 7, 8, 9, 13, 14, 15, 16, 17, 18] + : [1, 2, 3, 4, 5, 6, 7, 8, 12, 13, 14, 15, 17, 18]; + info = { ...info, widgets_values: indices.map((index, i) => ( + values[index] ?? this.widgets[i].value + )) }; + } + return configure.call(this, info); + }; + }, +}); diff --git a/src/Comfyui_Fd_Nodes/v2/rmbg_segment_v2_node.py b/src/Comfyui_Fd_Nodes/v2/rmbg_segment_v2_node.py index 69b4c77..5fbf189 100644 --- a/src/Comfyui_Fd_Nodes/v2/rmbg_segment_v2_node.py +++ b/src/Comfyui_Fd_Nodes/v2/rmbg_segment_v2_node.py @@ -5,7 +5,6 @@ from ..config import ( FD_BODY_SEGMENT_URL, FD_CLOTHES_SEGMENT_URL, FD_FASHION_SEGMENT_URL, - FD_RMBG_URL, RMBG_SERVICE_ENVS, ) from ..utils.image_payload import ( @@ -18,108 +17,14 @@ from ..zhiyi_rmbg_segment_node import ( BODY_CLASSES, CLOTHES_CLASSES, FASHION_CLASSES, - RMBG_IMAGE_ENCODINGS, - RMBG_IMAGE_ENCODING_CONFIG, + ZhiYiRMBGNode, _RmbgSegmentApiBase, ) from .env import resolve_service_env, service_env_options -class ZhiYiRMBGNodeV2(_RmbgSegmentApiBase): - """知衣 RMBG 2.0 背景去除节点 V2 - 支持灰度/开发/海外环境一键切换。""" - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "image": ("IMAGE",), - "env": (service_env_options(), { - "default": "开发环境", - "tooltip": "灰度/开发/海外环境使用内置地址（忽略 service_url 与 FD_* 环境变量）；自定义使用 service_url", - }), - "service_url": ("STRING", { - "default": FD_RMBG_URL or "http://10.1.0.230:8003/v1/rmbg", - "multiline": False, - "tooltip": "RMBG 2.0 背景去除接口地址，例如 http://host:8003/v1/rmbg；仅环境=自定义时使用", - }), - "process_res": ("INT", {"default": 1024, "min": 256, "max": 2048, "step": 64}), - "sensitivity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05}), - "mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1}), - "mask_offset": ("INT", {"default": 0, "min": -64, "max": 64, "step": 1}), - "invert_output": ("BOOLEAN", {"default": False}), - "refine_foreground": ("BOOLEAN", {"default": False}), - "background": (BACKGROUND_MODES, {"default": "Alpha"}), - "background_color": ("STRING", {"default": "#222222"}), - "return_image": ("BOOLEAN", {"default": True}), - "return_mask": ("BOOLEAN", {"default": True}), - "return_mask_image": ("BOOLEAN", {"default": False}), - "max_concurrency": ("INT", {"default": 1, "min": 1, "max": 16, "step": 1}), - "timeout": ("INT", {"default": 180, "min": 10, "max": 1200, "step": 10}), - }, - "optional": { - "health_check": ("BOOLEAN", {"default": False}), - "node_switch": ("INT", {"default": 0, "min": 0, "max": 1, "step": 1}), - "image_delivery": (IMAGE_DELIVERY_OPTIONS, image_delivery_input_config()), - "image_encoding": (RMBG_IMAGE_ENCODINGS, RMBG_IMAGE_ENCODING_CONFIG), - }, - } - - RETURN_TYPES = ("IMAGE", "MASK", "IMAGE", "JSON") - RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE", "INFO") - FUNCTION = "remove_background" - CATEGORY = "知衣/抠图" - OUTPUT_NODE = False - - def remove_background( - self, - image, - env, - service_url, - process_res, - sensitivity, - mask_blur, - mask_offset, - invert_output, - refine_foreground, - background, - background_color, - return_image, - return_mask, - return_mask_image, - max_concurrency, - timeout, - health_check=False, - node_switch=0, - image_delivery=IMAGE_DELIVERY_BASE64, - image_encoding: str = "png", - ): - resolved_base = resolve_service_env(env, RMBG_SERVICE_ENVS) - return self._execute( - image=image, - service_url=resolved_base or service_url, - endpoint="rmbg", - default_url=FD_RMBG_URL or "http://10.1.0.230:8003/v1/rmbg", - health_key="rmbg", - label="知衣RMBG2.0背景去除", - process_res=process_res, - mask_blur=mask_blur, - mask_offset=mask_offset, - invert_output=invert_output, - background=background, - background_color=background_color, - return_image=return_image, - return_mask=return_mask, - return_mask_image=return_mask_image, - max_concurrency=max_concurrency, - timeout=timeout, - health_check=health_check, - node_switch=node_switch, - service_url_preset=CUSTOM_SERVICE_URL_PRESET, - image_delivery=image_delivery, - image_encoding=image_encoding, - extra={"sensitivity": float(sensitivity), "refine_foreground": bool(refine_foreground)}, - info_extra={"sensitivity": float(sensitivity), "refine_foreground": bool(refine_foreground)}, - ) +class ZhiYiRMBGNodeV2(ZhiYiRMBGNode): + """RMBG V2 shares the deployment-controlled RMBG contract with V1.""" class ZhiYiClothesSegmentNodeV2(_RmbgSegmentApiBase): diff --git a/src/Comfyui_Fd_Nodes/zhiyi_rmbg_segment_node.py b/src/Comfyui_Fd_Nodes/zhiyi_rmbg_segment_node.py index 16e1c96..0555cc0 100644 --- a/src/Comfyui_Fd_Nodes/zhiyi_rmbg_segment_node.py +++ b/src/Comfyui_Fd_Nodes/zhiyi_rmbg_segment_node.py @@ -18,7 +18,9 @@ from .config import ( FD_CLOTHES_SEGMENT_URL, FD_DWPOSE_API_KEY, FD_FASHION_SEGMENT_URL, - FD_RMBG_URL, + FD_COMFY_RUNTIME_ENV, + RMBG_DEFAULT_IMAGE_ENCODING, + RMBG_RUNTIME_URL, RMBG_SERVICE_URL_PRESETS, ) from .utils.error_utils import ERROR_TIMEOUT, normalize_error_message @@ -44,8 +46,8 @@ DEFAULT_BODY_SEGMENT_URL = "http://10.1.0.230:8003/v1/segment/body" BACKGROUND_MODES = ["Alpha", "Color"] RMBG_IMAGE_ENCODINGS = ["png", "jpeg_q95"] RMBG_IMAGE_ENCODING_CONFIG = { - "default": "png", - "tooltip": "传输质量（单选）：png 无损；jpeg_q95 使用 JPEG Q95 推理副本节省流量，成图仍用原始 RGB。", + "default": RMBG_DEFAULT_IMAGE_ENCODING, + "tooltip": "传输质量：png 无损；jpeg_q95 节省流量。未精修时用原始 RGB 合成，精修时保留远端图片。", } @@ -530,6 +532,8 @@ class _RmbgSegmentApiBase: ) result = self._post_process(service_url, payload, timeout, label) + if payload_params["is_rmbg"] and not result.get("mask"): + raise RuntimeError("RMBG 响应缺少 mask") if payload_params["mask_only_response"]: mask = self._rmbg_mask_only_to_tensor(result, original_image.size) elif result.get("mask"): @@ -537,7 +541,15 @@ class _RmbgSegmentApiBase: else: mask = self._empty_mask_tensor(original_image.size) - if result.get("mask"): + if payload_params["extra"].get("refine_foreground") and payload_params["is_rmbg"]: + # The refined foreground RGB cannot be reconstructed from the source + mask. + processed_image = self._decode_png_data_url(result.get("processed_image"), "processed_image") + if processed_image.size != original_image.size: + processed_image = processed_image.resize(original_image.size, Image.Resampling.BILINEAR) + mode = "RGBA" if payload_params["background"] == "Alpha" else "RGB" + pixels = np.asarray(processed_image.convert(mode), dtype=np.float32) / 255.0 + result_image = torch.from_numpy(pixels.copy()).unsqueeze(0) + elif result.get("mask"): result_image = self._compose_output_image( image_tensor, mask, @@ -626,11 +638,11 @@ class _RmbgSegmentApiBase: RMBG_SERVICE_URL_PRESETS, ) final_service_url = _normalize_rmbg_segment_url(selected_service_url, default_url, endpoint) - mask_only_response = endpoint == "rmbg" + mask_only_response = endpoint == "rmbg" and not extra.get("refine_foreground", False) if image_encoding not in RMBG_IMAGE_ENCODINGS: raise ValueError(f"不支持的传输质量: {image_encoding}") - if mask_only_response: - return_image, return_mask, return_mask_image = False, True, False + if endpoint == "rmbg": + return_image, return_mask, return_mask_image = not mask_only_response, True, False image_tensors = self._expand_images(image) if not image_tensors: raise RuntimeError("未提供图片") @@ -649,6 +661,7 @@ class _RmbgSegmentApiBase: "return_mask": return_mask, "return_mask_image": return_mask_image, "mask_only_response": mask_only_response, + "is_rmbg": endpoint == "rmbg", "image_encoding": image_encoding, "extra": extra, } @@ -692,15 +705,10 @@ class ZhiYiRMBGNode(_RmbgSegmentApiBase): """知衣 RMBG 2.0 背景去除节点 - 纯 API 调用，不加载本地模型。""" @classmethod - def INPUT_TYPES(cls): + def INPUT_TYPES(cls) -> dict: return { "required": { "image": ("IMAGE",), - "service_url": ("STRING", { - "default": FD_RMBG_URL or DEFAULT_RMBG_URL, - "multiline": False, - "tooltip": "RMBG 2.0 背景去除接口地址，例如 http://host:8003/v1/rmbg", - }), "process_res": ("INT", {"default": 1024, "min": 256, "max": 2048, "step": 64}), "sensitivity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05}), "mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1}), @@ -709,19 +717,12 @@ class ZhiYiRMBGNode(_RmbgSegmentApiBase): "refine_foreground": ("BOOLEAN", {"default": False}), "background": (BACKGROUND_MODES, {"default": "Alpha"}), "background_color": ("STRING", {"default": "#222222"}), - "return_image": ("BOOLEAN", {"default": True}), - "return_mask": ("BOOLEAN", {"default": True}), - "return_mask_image": ("BOOLEAN", {"default": False}), "max_concurrency": ("INT", {"default": 1, "min": 1, "max": 16, "step": 1}), "timeout": ("INT", {"default": 180, "min": 10, "max": 1200, "step": 10}), }, "optional": { "health_check": ("BOOLEAN", {"default": False}), "node_switch": ("INT", {"default": 0, "min": 0, "max": 1, "step": 1}), - "service_url_preset": (service_url_preset_options(RMBG_SERVICE_URL_PRESETS), { - "default": CUSTOM_SERVICE_URL_PRESET, - "tooltip": "选择部署预设时忽略 service_url；选择自定义时使用 service_url 文本框", - }), "image_delivery": (IMAGE_DELIVERY_OPTIONS, image_delivery_input_config()), "image_encoding": (RMBG_IMAGE_ENCODINGS, RMBG_IMAGE_ENCODING_CONFIG), }, @@ -735,32 +736,29 @@ class ZhiYiRMBGNode(_RmbgSegmentApiBase): def remove_background( self, - image, - service_url, - process_res, - sensitivity, - mask_blur, - mask_offset, - invert_output, - refine_foreground, - background, - background_color, - return_image, - return_mask, - return_mask_image, - max_concurrency, - timeout, - health_check=False, - node_switch=0, - service_url_preset=None, - image_delivery=IMAGE_DELIVERY_BASE64, - image_encoding: str = "png", - ): + image: torch.Tensor, + process_res: int, + sensitivity: float, + mask_blur: int, + mask_offset: int, + invert_output: bool, + refine_foreground: bool, + background: str, + background_color: str, + max_concurrency: int, + timeout: int, + health_check: bool = False, + node_switch: int = 0, + image_delivery: str = IMAGE_DELIVERY_BASE64, + image_encoding: str = RMBG_DEFAULT_IMAGE_ENCODING, + **legacy_inputs: object, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, str]: + # Old API graphs may still contain routing/return fields; deployment wins. return self._execute( image=image, - service_url=service_url, + service_url=RMBG_RUNTIME_URL, endpoint="rmbg", - default_url=FD_RMBG_URL or DEFAULT_RMBG_URL, + default_url=RMBG_RUNTIME_URL, health_key="rmbg", label="知衣RMBG2.0背景去除", process_res=process_res, @@ -769,18 +767,19 @@ class ZhiYiRMBGNode(_RmbgSegmentApiBase): invert_output=invert_output, background=background, background_color=background_color, - return_image=return_image, - return_mask=return_mask, - return_mask_image=return_mask_image, + return_image=bool(refine_foreground), + return_mask=True, + return_mask_image=False, max_concurrency=max_concurrency, timeout=timeout, health_check=health_check, node_switch=node_switch, - service_url_preset=service_url_preset, + service_url_preset=CUSTOM_SERVICE_URL_PRESET, image_delivery=image_delivery, image_encoding=image_encoding, extra={"sensitivity": float(sensitivity), "refine_foreground": bool(refine_foreground)}, - info_extra={"sensitivity": float(sensitivity), "refine_foreground": bool(refine_foreground)}, + info_extra={"sensitivity": float(sensitivity), "refine_foreground": bool(refine_foreground), + "runtime_env": FD_COMFY_RUNTIME_ENV}, ) diff --git a/tests/test_rmbg_mask_only.py b/tests/test_rmbg_mask_only.py index d0d4f45..0a4bc1a 100644 --- a/tests/test_rmbg_mask_only.py +++ b/tests/test_rmbg_mask_only.py @@ -105,6 +105,7 @@ def test_node_quality_ignores_old_environment_switch(monkeypatch, node_cls: type monkeypatch.setattr(module, "tensor_to_image_payload", lambda *args: "data:image/png;base64,input") monkeypatch.setattr(module, "tensor_to_jpeg_q95_payload", lambda *args: "data:image/jpeg;base64,input") node = node_cls() + monkeypatch.setattr(module, "RMBG_RUNTIME_URL", "http://domestic.example/v1/rmbg") def post(url: str, payload: dict, timeout: int, label: str) -> dict: assert url == "http://domestic.example/v1/rmbg" diff --git a/tests/test_rmbg_runtime.py b/tests/test_rmbg_runtime.py new file mode 100644 index 0000000..53fc0c2 --- /dev/null +++ b/tests/test_rmbg_runtime.py @@ -0,0 +1,82 @@ +"""Deployment routing, legacy graph isolation and refined pixel preservation.""" + +import base64 +import io +import json +import os +from pathlib import Path +import subprocess +import sys + +import numpy as np +from PIL import Image +import pytest +import torch + +from src.Comfyui_Fd_Nodes import zhiyi_rmbg_segment_node as module +from src.Comfyui_Fd_Nodes.v2.rmbg_segment_v2_node import ZhiYiRMBGNodeV2 +from tests.test_rmbg_mask_only import call_node, mask_response + + +@pytest.mark.parametrize("runtime,quality,host", [ + ("dev", "png", "10.1.0.230:8003"), + ("domestic", "png", "model-api-dwpose-svc.online-server-gray:8001"), + ("overseas", "jpeg_q95", "overseas.example/model-api-dwpose"), +]) +@pytest.mark.parametrize("override", [False, True]) +def test_runtime_config(runtime: str, quality: str, host: str, override: bool) -> None: + env = {**os.environ, "FD_COMFY_RUNTIME_ENV": runtime, "FD_OVERSEAS_GPU_GATEWAY": "http://overseas.example"} + env.pop("FD_RMBG_URL", None) + if override: + env["FD_RMBG_URL"] = "http://configured.example/v1/rmbg" + config = Path(module.__file__).with_name("config.py") + code = "import runpy,json,sys; c=runpy.run_path(sys.argv[1]); print(json.dumps([c['RMBG_RUNTIME_URL'],c['RMBG_DEFAULT_IMAGE_ENCODING']]))" + result = subprocess.run([sys.executable, "-c", code, str(config)], env=env, capture_output=True, text=True, check=True) + assert json.loads(result.stdout) == [f"http://{host}/v1/rmbg" if not override else env["FD_RMBG_URL"], quality] + + +def test_invalid_runtime_fails_on_load() -> None: + env = {**os.environ, "FD_COMFY_RUNTIME_ENV": "typo"} + result = subprocess.run([sys.executable, str(Path(module.__file__).with_name("config.py"))], + env=env, capture_output=True, text=True) + assert result.returncode != 0 and "FD_COMFY_RUNTIME_ENV must be" in result.stderr + + +@pytest.mark.parametrize("node_cls", [module.ZhiYiRMBGNode, ZhiYiRMBGNodeV2]) +@pytest.mark.parametrize("background", ["Alpha", "Color"]) +def test_refine_preserves_server_pixels(monkeypatch, node_cls: type, background: str) -> None: + node = node_cls() + mode = "RGBA" if background == "Alpha" else "RGB" + color = (37, 111, 207, 128) if mode == "RGBA" else (37, 111, 207) + refined = Image.new(mode, (5, 3), color) + buffer = io.BytesIO() + refined.save(buffer, format="PNG") + response = mask_response() + response["processed_image"] = base64.b64encode(buffer.getvalue()).decode() + + def post(url: str, payload: dict, timeout: int, label: str) -> dict: + assert (payload["return_image"], payload["return_mask"], payload["return_mask_image"]) == (True, True, False) + assert payload["refine_foreground"] is True + return response + + monkeypatch.setattr(node, "_post_process", post) + image = torch.zeros(1, 3, 5, 3) + output, mask, preview, text = call_node(node, image, refine_foreground=True, background=background, + return_image=False, return_mask=False) + assert torch.equal(output, torch.from_numpy(np.asarray(refined).copy().astype(np.float32) / 255).unsqueeze(0)) + assert torch.equal(preview, mask.unsqueeze(-1).expand(-1, -1, -1, 3)) + assert json.loads(text)["mask_only_response"] is False + + +def test_refine_missing_image_fails(monkeypatch) -> None: + node = ZhiYiRMBGNodeV2() + monkeypatch.setattr(node, "_post_process", lambda *args: mask_response()) + with pytest.raises(RuntimeError, match="processed_image"): + call_node(node, torch.zeros(1, 3, 5, 3), refine_foreground=True) + + +@pytest.mark.parametrize("node_cls", [module.ZhiYiRMBGNode, ZhiYiRMBGNodeV2]) +def test_retired_controls_are_not_node_inputs(node_cls: type) -> None: + fields = node_cls.INPUT_TYPES() + assert not {"env", "service_url", "service_url_preset", "return_image", "return_mask", "return_mask_image"} & ( + fields["required"].keys() | fields["optional"].keys()) diff --git a/tests/test_rmbg_segment_node.py b/tests/test_rmbg_segment_node.py index 6b49ae4..9eb4d6a 100644 --- a/tests/test_rmbg_segment_node.py +++ b/tests/test_rmbg_segment_node.py @@ -31,6 +31,11 @@ from src.Comfyui_Fd_Nodes.zhiyi_rmbg_segment_node import ( ) +@pytest.fixture(autouse=True) +def rmbg_deployment_url(monkeypatch): + monkeypatch.setattr(rmbg_module, "RMBG_RUNTIME_URL", "https://example.com/api/v1/rmbg") + + def _image_data_url(size=(2, 2), color=(0, 255, 0), mode="RGB"): image = Image.new(mode, size, color) buffer = io.BytesIO() @@ -101,11 +106,11 @@ def test_rmbg_and_segment_node_metadata(): fashion_inputs = ZhiYiFashionSegmentNode.INPUT_TYPES() body_inputs = ZhiYiBodySegmentNode.INPUT_TYPES() - assert rmbg_inputs["required"]["service_url"][1]["default"].endswith("/v1/rmbg") + assert "service_url" not in rmbg_inputs["required"] assert rmbg_inputs["required"]["process_res"][1]["default"] == 1024 assert rmbg_inputs["required"]["sensitivity"][1]["default"] == 1.0 assert rmbg_inputs["required"]["refine_foreground"][1]["default"] is False - assert rmbg_inputs["required"]["return_mask_image"][1]["default"] is False + assert not {"return_image", "return_mask", "return_mask_image"} & rmbg_inputs["required"].keys() assert clothes_inputs["required"]["service_url"][1]["default"].endswith("/v1/segment/clothes") assert fashion_inputs["required"]["service_url"][1]["default"].endswith("/v1/segment/fashion") assert body_inputs["required"]["service_url"][1]["default"].endswith("/v1/segment/body") @@ -141,6 +146,8 @@ def test_rmbg_and_segment_node_metadata(): assert node_cls.RETURN_TYPES == ("IMAGE", "MASK", "IMAGE", "JSON") assert node_cls.RETURN_NAMES == ("IMAGE", "MASK", "MASK_IMAGE", "INFO") input_types = node_cls.INPUT_TYPES() + if node_cls is ZhiYiRMBGNode: + continue assert list(input_types["optional"])[2:4] == ["service_url_preset", "image_delivery"] preset_options, preset_config = input_types["optional"]["service_url_preset"] assert preset_options == [CUSTOM_SERVICE_URL_PRESET, "10.1.0.230"] @@ -257,6 +264,7 @@ def test_parse_classes_preserves_fashion_comma_names(): def test_rmbg_posts_full_payload_and_decodes_outputs(monkeypatch): + monkeypatch.setattr(rmbg_module, "RMBG_RUNTIME_URL", "http://example.com:8003/v1/rmbg") node = ZhiYiRMBGNode() images = torch.zeros((1, 2, 2, 3), dtype=torch.float32) calls = [] @@ -276,7 +284,7 @@ def test_rmbg_posts_full_payload_and_decodes_outputs(monkeypatch): assert payload["refine_foreground"] is True assert payload["background"] == "Color" assert payload["background_color"] == "#ffffff" - assert payload["return_image"] is False + assert payload["return_image"] is True assert payload["return_mask"] is True assert payload["return_mask_image"] is False assert payload["image_format"] == "png_base64" @@ -284,7 +292,7 @@ def test_rmbg_posts_full_payload_and_decodes_outputs(monkeypatch): request_id="rmbg-1", classes=[], preprocessor="RMBGPreprocessor", - processed_image=None, + processed_image=_image_data_url(color=(0, 255, 0)), mask_image=None, ) data["model_name"] = "RMBG-2.0" @@ -319,8 +327,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) - # Compose output and grayscale preview locally from the returned mask. - assert torch.equal(result_image[0, 0, 0], torch.tensor([0.0, 0.0, 0.0])) + # Preserve the refined server RGB; only the preview is generated locally. + assert torch.equal(result_image[0, 0, 0], torch.tensor([0.0, 1.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, 1.0, 1.0])) assert info["service_url"] == "http://example.com:8003/v1/rmbg" @@ -543,7 +551,7 @@ def test_image_delivery_appended_and_default_preserves_signature(node_cls, endpo inputs = node_cls.INPUT_TYPES() is_v2 = "env" in inputs["required"] old_optional = ["health_check", "node_switch"] - if not is_v2: + if not is_v2 and endpoint != "rmbg": old_optional.append("service_url_preset") assert list(inputs["optional"]) == old_optional + ["image_delivery"] + (["image_encoding"] if endpoint == "rmbg" else []) options, config = inputs["optional"]["image_delivery"] @@ -551,7 +559,7 @@ def test_image_delivery_appended_and_default_preserves_signature(node_cls, endpo assert config == image_payload.image_delivery_input_config() method = getattr(node_cls, node_cls.FUNCTION) parameters = inspect.signature(method).parameters - assert list(parameters) == ["self", *inputs["required"], *inputs["optional"]] + assert list(parameters) == ["self", *inputs["required"], *inputs["optional"], *(["legacy_inputs"] if endpoint == "rmbg" else [])] assert parameters["image_delivery"].default == "base64" diff --git a/tests/test_v2_nodes.py b/tests/test_v2_nodes.py index a46ba1c..1b24160 100644 --- a/tests/test_v2_nodes.py +++ b/tests/test_v2_nodes.py @@ -201,7 +201,9 @@ def test_input_types_env_first_service_url_second(): for cls in V2_CLASSES: input_types = cls.INPUT_TYPES() required_keys = list(input_types["required"]) - if cls is ZhiYiSAM2SegmentNodeV2: + if cls is ZhiYiRMBGNodeV2: + assert required_keys[1] == "process_res" + elif cls is ZhiYiSAM2SegmentNodeV2: assert required_keys[0] == "images" assert required_keys[1] == "bboxes" assert required_keys[2] == "env" @@ -220,7 +222,9 @@ def test_input_types_business_params_match_old(): for v2_cls, old_cls in zip(V2_CLASSES, OLD_CLASSES): v2_keys = list(v2_cls.INPUT_TYPES()["required"]) old_keys = list(old_cls.INPUT_TYPES()["required"]) - if v2_cls is ZhiYiSAM2SegmentNodeV2: + if v2_cls is ZhiYiRMBGNodeV2: + assert v2_keys == old_keys + elif v2_cls is ZhiYiSAM2SegmentNodeV2: assert v2_keys[4:] == old_keys[3:], v2_cls.__name__ else: filtered = [key for key in v2_keys[3:] if key not in switch_filters.get(v2_cls, set())] @@ -235,6 +239,9 @@ def test_env_options_shape(): def test_env_combo_is_string_options(): for cls in V2_CLASSES: + if cls is ZhiYiRMBGNodeV2: + assert "env" not in cls.INPUT_TYPES()["required"] + continue options, config = cls.INPUT_TYPES()["required"]["env"] assert all(isinstance(option, str) for option in options) @@ -243,7 +250,7 @@ def test_env_combo_is_string_options(): def test_env_defaults(): - for cls in [ZhiYiRMBGNodeV2, ZhiYiClothesSegmentNodeV2, ZhiYiFashionSegmentNodeV2, ZhiYiBodySegmentNodeV2]: + for cls in [ZhiYiClothesSegmentNodeV2, ZhiYiFashionSegmentNodeV2, ZhiYiBodySegmentNodeV2]: assert cls.INPUT_TYPES()["required"]["env"][1]["default"] == "开发环境" for cls in [ZhiYiDWPoseDetectNodeV2, ZhiYiLineArtPreprocessorNodeV2, ZhiYiDepthAnythingV2PreprocessorNodeV2, ZhiYiSAM2SegmentNodeV2]: @@ -353,7 +360,8 @@ def test_rmbg_v2_dev_env_url(monkeypatch): def fake_request(method, url, timeout, **kwargs): calls.append(url) - return DummyResponse(data=_segment_response()) + return DummyResponse(data={**_segment_response(), "processed_image": None, "mask_image": None, + "inference_size": {"width": 2, "height": 2}, "output_size": {"width": 2, "height": 2}}) monkeypatch.setattr(rmbg_module.requests, "request", fake_request) @@ -695,33 +703,18 @@ def test_sam2_v2_three_envs_post_urls(monkeypatch): ] -def test_rmbg_v2_invalid_env_raises(monkeypatch): +def test_rmbg_v2_legacy_env_cannot_override_deployment(monkeypatch): node = ZhiYiRMBGNodeV2() + captured = {} - def fake_request(method, url, timeout, **kwargs): - raise AssertionError("不应发出请求") + def execute(**kwargs): + captured.update(kwargs) + return None - monkeypatch.setattr(rmbg_module.requests, "request", fake_request) - - with pytest.raises(RuntimeError, match="未知的 env: 不存在"): - node.remove_background( - image=torch.zeros((1, 2, 2, 3), dtype=torch.float32), - env="不存在", - service_url="", - process_res=1024, - sensitivity=1.0, - mask_blur=0, - mask_offset=0, - invert_output=False, - refine_foreground=False, - background="Alpha", - background_color="#222222", - return_image=True, - return_mask=True, - return_mask_image=False, - max_concurrency=1, - timeout=30, - ) + monkeypatch.setattr(node, "_execute", execute) + kwargs = {name: spec[1]["default"] for name, spec in node.INPUT_TYPES()["required"].items() if name != "image"} + node.remove_background(image=torch.zeros(1, 2, 2, 3), env="不存在", service_url="http://wrong", **kwargs) + assert captured["service_url"] == rmbg_module.RMBG_RUNTIME_URL # ---------- 9. 老节点回归保护 ---------- @@ -742,6 +735,8 @@ def test_old_controlnet_payload_still_returns_image_true(): def test_old_nodes_unchanged_input_types(): for old_cls in OLD_CLASSES: + if old_cls is ZhiYiRMBGNode: + continue input_types = old_cls.INPUT_TYPES() assert "env" not in input_types["required"] assert list(input_types["required"])[1] == "bboxes" or list(input_types["required"])[1] == "service_url" @@ -773,7 +768,8 @@ def test_v2_node_switch_short_circuit(monkeypatch): timeout=30, node_switch=1, ) - assert torch.equal(result_image, image) + assert result_image.shape == (1, 2, 2, 4) + assert torch.count_nonzero(result_image) == 0 assert json.loads(info_text)["skipped"] is True dwpose_node = ZhiYiDWPoseDetectNodeV2()