nodes/SeanBRVFX-ComfyUI-CorridorKey
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Initial ComfyUI-CorridorKey custom node

Browse Source
This commit is contained in:
SeanBRVFX
2026-03-02 10:00:09 +01:00
commit 121bec58b0
13 changed files with 1497 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+9
View File
@@ -0,0 +1,9 @@
__pycache__/
*.py[cod]
.pytest_cache/
.mypy_cache/
.ruff_cache/
.venv/
dist/
build/
*.egg-info/
LICENSE
+25
View File
@@ -0,0 +1,25 @@
ComfyUI-CorridorKey License Notice
This custom node is intended to follow the licensing expectations stated by the
upstream CorridorKey repository by nikopueringer:
https://github.com/nikopueringer/CorridorKey
Summary of the upstream terms as described in the upstream README:
1. The upstream project is not presented as a standard standalone SPDX license
file. The README describes the release as effectively a variation of
CC BY-NC-SA 4.0.
2. Repackaging and reselling the tool is not permitted.
3. Offering paid API inference based on the model is not permitted without a
separate agreement from the upstream author.
4. Released variations or improvements should remain free and open source.
5. Future forks or releases should keep the "Corridor Key" name.
This file is a convenience summary only. The authoritative reference for these
terms is the upstream README section titled:
"CorridorKey Licensing and Permissions"
If the upstream author publishes a formal license file or updated terms, those
terms should be reviewed and this file should be updated accordingly.
README.md
+468
View File
@@ -0,0 +1,468 @@
# ComfyUI-CorridorKey
## 1) Project name + one line summary
`ComfyUI-CorridorKey` is a ComfyUI custom node package that exposes a native `CorridorKey` node for CorridorKey-style coarse-mask refinement inside ComfyUI.
## 2) Problem statement and goals (what it does, what it does NOT do)
`ComfyUI-CorridorKey` adapts the intent of the upstream `CorridorKey` project for ComfyUI users who already build masks inside ComfyUI. Instead of running a separate command-line application and instead of bundling extra mask-generation tools such as `VideoMaMa` or `GVM`, this node accepts:
- an `IMAGE` input
- a coarse `MASK` input produced by other ComfyUI nodes
It then performs an edge-aware matte refinement pass and returns:
- a refined alpha `MASK`
- a selected preview `IMAGE` based on the requested output mode
- dedicated `IMAGE` outputs for foreground, processed, and comp previews
Goals:
- Keep the integration native to ComfyUI's standard custom-node format.
- Reuse existing ComfyUI nodes for coarse mask creation.
- Keep runtime dependencies small and compatible with a typical ComfyUI Python environment.
- Provide a deterministic, testable baseline refinement pipeline that is safe to run locally.
- Track upstream `CorridorKey` changes safely and surface newer verified commits without blindly overwriting local custom-node code.
Non-goals:
- This project does not reimplement the full upstream research model or ship upstream CorridorKey checkpoints.
- This project does not recreate `VideoMaMa`, `GVM`, or any separate mask-proposal pipeline.
- This project does not download models, call external services, or require network access at runtime.
- This project does not write output files to disk; it only returns ComfyUI node outputs.
- This project does not auto-overwrite its own code from the upstream `CorridorKey` repository, because that repository is not the same codebase as this ComfyUI custom node.
## 3) Features (current + clearly marked planned)
Current features:
- One ComfyUI node: `CorridorKey`
- Accepts batched `IMAGE` tensors and `MASK` tensors
- Broadcasts a single mask across an image batch when needed
- Validates shapes, value ranges, and numeric settings
- Exposes upstream-inspired controls for:
- `Gamma Space`
- `Despill Strength`
- `Auto-Despeckle`
- `Despeckle Size`
- `Refiner Strength`
- output mode selection (`FG`, `Matte`, `Processed`, `Comp`)
- Performs deterministic edge-aware mask smoothing, spill suppression, and comp preview generation
- Returns refined alpha mask plus selected output, foreground, processed, and comp preview images
- On module load, performs a best-effort background check for newer upstream commits that have no failing GitHub check-runs
- Includes pytest coverage for the core processor
Planned features:
- Planned: optional backend that can call a true upstream CorridorKey-compatible runtime if the user provides one
- Planned: trimap output for downstream compositing workflows
- Planned: preset profiles tuned for hair, glass, and semi-transparent edges
- Planned: optional preview compositing against solid-color backgrounds
## 4) Requirements (Python version, OS, dependencies)
- Python: 3.10 or newer
- OS: Windows, Linux, or macOS where ComfyUI runs
- Runtime dependencies:
- `torch` (normally already present in the ComfyUI Python environment)
- `numpy`
- `Pillow`
- Development dependencies:
- `pytest`
- `ruff`
- `black`
- `mypy`
## 5) Quick start (venv, install, run)
This repository is meant to live under ComfyUI's `custom_nodes` directory.
Recommended ComfyUI installation:
1. Clone this repository into `ComfyUI/custom_nodes/`.
2. Install the runtime dependencies into the same Python environment ComfyUI uses.
3. Restart ComfyUI.
Typical install flow:
```powershell
cd ComfyUI\custom_nodes
git clone <repo-url> ComfyUI-CorridorKey
cd ComfyUI-CorridorKey
python -m pip install -r requirements.txt
```
After that, restart ComfyUI and the `CorridorKey` node should appear.
If you are using ComfyUI Portable on Windows and want to be explicit about the embedded Python:
```powershell
cd ComfyUI_windows_portable\ComfyUI\custom_nodes
git clone <repo-url> ComfyUI-CorridorKey
..\..\python_embeded\python.exe -m pip install -r .\ComfyUI-CorridorKey\requirements.txt
```
Manual install also works:
1. Copy the `ComfyUI-CorridorKey` folder into `ComfyUI/custom_nodes/`.
2. Run `python -m pip install -r requirements.txt` from inside that folder.
3. Restart ComfyUI.
Development install:
```powershell
python -m venv .venv
.venv\Scripts\Activate.ps1
python -m pip install --upgrade pip
python -m pip install -r requirements.txt
python -m pip install -e .[dev]
```
Restart ComfyUI after installation. The node will appear under the `CorridorKey` category.
For a normal ComfyUI install, `requirements.txt` is the primary runtime dependency file. `pyproject.toml` is kept for editable installs and development tooling.
## 6) Configuration (env vars, config files, defaults)
The node does not require config files.
Runtime configuration is passed through node inputs:
- `gamma_space`: `sRGB` or `Linear`, matching the upstream input interpretation concept
- `feather_radius`: integer blur radius used to soften the coarse mask before refinement
- `edge_focus`: float multiplier that strengthens image-edge preservation
- `threshold`: float cutoff used to center the matte transition
- `preserve_core`: float band that locks obvious foreground/background values
- `despill_strength`: normalized float from `0.0` to `1.0` that reduces green spill on semi-transparent foreground edges
- `refiner_strength`: float multiplier that strengthens the edge-aware matte correction pass
- `auto_despeckle`: enables a small morphology-based cleanup pass
- `despeckle_size`: integer control for the minimum speckle scale to suppress
- `output_mode`: selects the primary preview output as `FG`, `Matte`, `Processed`, or `Comp`
Defaults:
- `gamma_space = sRGB`
- `feather_radius = 4`
- `edge_focus = 1.5`
- `threshold = 0.5`
- `preserve_core = 0.2`
- `despill_strength = 1.0`
- `refiner_strength = 1.0`
- `auto_despeckle = enabled`
- `despeckle_size = 400`
- `output_mode = Processed`
Implementation note:
- The upstream app uses model-driven processing. This ComfyUI node mirrors the same user-facing control surface where practical, but the internal math is a deterministic local approximation built around the provided image and coarse mask.
- The upstream CLI uses a numeric prompt for `Auto-Despeckle Size` with a default of `400` pixels. This node exposes the same value as an integer input rather than forcing a fixed dropdown.
Environment variables for upstream tracking:
- `CORRIDORKEY_AUTO_CHECK_UPSTREAM`
- Default: `1`
- Set to `0` to disable the background upstream check
- `CORRIDORKEY_UPSTREAM_TIMEOUT_SECONDS`
- Default: `3.0`
- Per-request timeout for GitHub API calls
- `CORRIDORKEY_UPSTREAM_CHECK_DEPTH`
- Default: `15`
- How many recent upstream commits to scan when looking for the newest verified commit
The currently reviewed upstream head SHA and its last observed check state are pinned in `corridor_key/upstream_sync.py`.
## 7) Usage (examples, CLI/API examples if relevant)
Recommended ComfyUI flow:
1. Load or generate an image.
2. Create a coarse mask using existing ComfyUI nodes such as RMBG, SAM, Segment Anything, Florence-based tools, or manual masking nodes.
3. Feed the `IMAGE` and coarse `MASK` into `CorridorKey`.
4. Choose `FG`, `Matte`, `Processed`, or `Comp` as the node's selected preview output.
5. Use the refined `MASK` for compositing, inpainting, EXR assembly, or export.
Node inputs:
- `image`: `IMAGE`
- `mask`: `MASK`
- `gamma_space`: `COMBO`
- `feather_radius`: `INT`
- `edge_focus`: `FLOAT`
- `threshold`: `FLOAT`
- `preserve_core`: `FLOAT`
- `despill_strength`: `FLOAT`
- `refiner_strength`: `FLOAT`
- `auto_despeckle`: `COMBO`
- `despeckle_size`: `INT`
- `output_mode`: `COMBO`
Node outputs:
- `refined_mask`: `MASK`
- `selected_output`: `IMAGE`
- `foreground`: `IMAGE`
- `processed`: `IMAGE`
- `comp`: `IMAGE`
Example workflow idea:
- `Load Image` -> `RMBG` (or another segmentation node) -> `CorridorKey` -> `Preview Image`
- For EXR-oriented workflows, use the returned `foreground` and `refined_mask` separately and combine or export them with other ComfyUI nodes such as `cocotools_io`, instead of expecting this node to write EXR files directly.
There is no standalone CLI in this custom node package by design.
Upstream tracking behavior:
- On import, the package starts a small background check.
- It queries the upstream `nikopueringer/CorridorKey` GitHub repo.
- It looks for the newest recent commit with at least one successful check-run and no failing check-runs.
- If such a commit is newer than the currently pinned verified upstream baseline, it logs that an upstream update is available for manual review.
- It does not automatically rewrite local files from upstream.
## 8) Architecture overview
### High-level diagram in ASCII
```text
IMAGE + coarse MASK
|
v
input validation
|
v
mask normalization
|
v
optional gamma handling
|
v
image luminance + edge map
|
v
refiner + despeckle
|
v
refined alpha MASK
|
v
despill + output assembly
| | |
v v v
FG Processed Comp
```
### Key modules and responsibilities
- `nodes.py`
- ComfyUI-facing node class
- Defines `INPUT_TYPES`, return types, and node metadata
- Converts node inputs into the internal processor call
- `corridor_key/config.py`
- Dataclass for validated runtime settings
- `corridor_key/tensor_ops.py`
- Shared tensor validation, mask normalization, color helpers, and small image-processing helpers
- `corridor_key/processor.py`
- Pure processing pipeline for edge-aware matte refinement, despill, despeckle, and output assembly
- `corridor_key/upstream_sync.py`
- Best-effort GitHub API integration for verified-upstream commit checks
- `tests/test_processor.py`
- Unit tests for shape handling, validation, and output behavior
### Data flow / request flow
1. ComfyUI invokes `CorridorKey.run`.
2. The node builds a `CorridorKeySettings` object from numeric inputs.
3. The processor validates the image and mask tensors.
4. A normalized mask batch is created and aligned to the image batch.
5. The processor optionally adapts handling for `Gamma Space`.
6. The processor derives luminance, computes an edge map, and applies the refiner pass.
7. Optional despeckle cleanup is applied when enabled.
8. The processor builds foreground, processed, and comp previews, including optional despill.
9. The processor returns:
- refined mask
- selected output
- foreground
- processed
- comp
10. ComfyUI receives the outputs without any file-system side effects.
11. In parallel, an optional background task may query GitHub and log if a newer verified upstream commit is available.
## 9) Repository structure (tree)
```text
ComfyUI-CorridorKey/
|-- README.md
|-- LICENSE
|-- requirements.txt
|-- pyproject.toml
|-- __init__.py
|-- nodes.py
|-- corridor_key/
| |-- __init__.py
| |-- config.py
| |-- tensor_ops.py
| |-- processor.py
| `-- upstream_sync.py
`-- tests/
`-- test_processor.py
```
## 10) Development
### local dev loop
```powershell
python -m pip install -e .[dev]
pytest -q
ruff check .
black --check .
mypy corridor_key
```
When changing behavior:
1. Update `README.md` first if the behavior, architecture, setup, or workflow changes.
2. Keep the implementation aligned with the documented repository structure.
3. Add or update pytest coverage for the changed logic.
4. Restart ComfyUI and re-test the node in a simple image-plus-mask workflow.
### testing strategy (pytest)
- Test pure processing code in `corridor_key/processor.py`
- Avoid ComfyUI runtime dependencies in the test suite where possible
- Cover:
- mask broadcasting
- parameter validation
- output shapes
- value range clamping
- output mode selection
- despeckle and despill edge cases
- verified-upstream selection policy for mocked GitHub metadata
### lint/format rules (ruff + black or ruff format)
- `ruff check .` for linting
- `black .` for formatting
- Keep imports explicit and functions small
- Prefer readable tensor code over compact but opaque expressions
### type checking (mypy or pyright if useful)
- `mypy corridor_key`
- Use standard type hints for public helpers and processor methods
- Keep the ComfyUI adapter lightweight and typed where practical
## 11) Error handling and logging (logging module, structured messages)
- Internal validation raises `ValueError` with specific, user-readable messages
- The ComfyUI node catches nothing silently; invalid inputs should fail clearly
- Logging is intentionally minimal:
- a module-level logger may emit debug-level messages in the processor
- a background upstream checker may emit info-level status lines
- no print-based progress logging
- GitHub API checks use explicit short timeouts and fail closed
## 12) Security notes (inputs, secrets, safe defaults)
- No `eval`, `exec`, or unsafe deserialization
- No subprocess execution
- No model downloads or arbitrary network access during image processing
- No secrets are required
- All user-provided numeric inputs are validated and clamped to safe ranges
- File paths are not accepted as node inputs, so the node does not read or write arbitrary files
- The processor only operates on in-memory tensors returned by ComfyUI
- The optional upstream checker uses only GitHub's API with a short timeout, does not require tokens, and never executes remote code
## 13) Deployment (if applicable)
Deployment is simply installation as a ComfyUI custom node:
1. Place this folder under `ComfyUI/custom_nodes/`.
2. Install the package into the same Python environment used by ComfyUI.
3. Restart ComfyUI.
There is no separate server, container, or packaged service in this repository.
## 14) Troubleshooting and FAQ
Q: Why does this node require a mask input?
A: This package intentionally reuses existing ComfyUI segmentation and masking nodes instead of duplicating them.
Q: Why does it not match the full upstream CorridorKey model exactly?
A: This implementation provides a native, deterministic refinement pass inspired by the same workflow goal, without bundling upstream research weights or optional external pipelines.
Q: My mask looks too soft.
A: Lower `feather_radius`, lower `preserve_core`, raise `edge_focus`, or raise `refiner_strength`.
Q: My mask looks too harsh.
A: Raise `feather_radius` slightly or lower `threshold`.
Q: Why is `Processed` not writing an EXR with alpha by itself?
A: ComfyUI nodes exchange in-memory tensors. This node returns separate `foreground` and `refined_mask` outputs so you can assemble or export EXR data elsewhere in the workflow.
Q: Why is the default `despeckle_size` so large?
A: The upstream `CorridorKey` CLI defaults `Auto-Despeckle Size` to `400` pixels. On small images, you may need to lower that value or disable auto-despeckle.
Q: Does this node auto-update itself from GitHub?
A: It automatically checks for newer upstream commits that appear verified, but it does not self-modify local code. The upstream repo is a different project, so updates should still be reviewed and ported intentionally.
## 15) Roadmap
- Add an optional pluggable backend for a true upstream-compatible inference engine
- Add trimap and confidence outputs
- Add more regression tests around semi-transparent regions
- Add a true RGBA or EXR-friendly packed output path if the surrounding ComfyUI ecosystem provides a stable tensor contract for it
- Add an example ComfyUI workflow JSON once the node behavior stabilizes
## 16) Contributing (branching, commit style, PR checklist)
Branching:
- Use short-lived feature branches from `main`
- Keep one focused change per branch
Commit style:
- Prefer imperative commit messages
- Example: `Add edge-aware CorridorKey refinement node`
PR checklist:
- Update `README.md` first when behavior or structure changes
- Run `pytest -q`
- Run `ruff check .`
- Run `black --check .`
- Run `mypy corridor_key`
- Confirm the node loads in ComfyUI
- Keep the public node contract backward compatible unless documented
## 17) License placeholder
This package is based on the workflow and licensing expectations stated by the upstream `nikopueringer/CorridorKey` repository.
Upstream license status:
- The upstream repository does not present a standard standalone SPDX license file in the root.
- Its README includes a custom license statement described by the author as effectively a variation of `CC BY-NC-SA 4.0`.
- The upstream README also states additional restrictions, including:
- no repackaging and resale of the tool
- no paid API inference service using the model without a separate agreement
- released variations or improvements should remain free and open source
- future forks or releases should keep the `Corridor Key` name
For this custom node, treat the upstream repository's licensing terms as the governing reference unless the repository owner explicitly relicenses this work.
This repository also includes a local `LICENSE` file that summarizes the same upstream-derived terms for convenience.
Source reference:
- Upstream repository: `https://github.com/nikopueringer/CorridorKey`
- Upstream licensing section: `README.md`, section `CorridorKey Licensing and Permissions`
__init__.py
+33
View File
@@ -0,0 +1,33 @@
"""
@author: local
@title: ComfyUI-CorridorKey
@nickname: CorridorKey
@description: ComfyUI node for CorridorKey-style coarse mask refinement.
"""
from __future__ import annotations
import sys
from pathlib import Path
try:
from .corridor_key import schedule_upstream_check
from .nodes import CorridorKey
except ImportError:
package_dir = Path(__file__).resolve().parent
if str(package_dir) not in sys.path:
sys.path.insert(0, str(package_dir))
from corridor_key import schedule_upstream_check
from nodes import CorridorKey
NODE_CLASS_MAPPINGS = {
"CorridorKey": CorridorKey,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"CorridorKey": "CorridorKey",
}
schedule_upstream_check()
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
corridor_key/__init__.py
+25
View File
@@ -0,0 +1,25 @@
from .config import CorridorKeySettings
from .processor import CorridorKeyProcessor
from .upstream_sync import (
SYNCED_UPSTREAM_HEAD_CHECK_CONCLUSIONS,
SYNCED_UPSTREAM_HEAD_DATE,
SYNCED_UPSTREAM_HEAD_MESSAGE,
SYNCED_UPSTREAM_HEAD_SHA,
UpstreamCommitRecord,
is_verified_check_conclusions,
schedule_upstream_check,
select_latest_verified_commit,
)
__all__ = [
"CorridorKeyProcessor",
"CorridorKeySettings",
"SYNCED_UPSTREAM_HEAD_CHECK_CONCLUSIONS",
"SYNCED_UPSTREAM_HEAD_DATE",
"SYNCED_UPSTREAM_HEAD_MESSAGE",
"SYNCED_UPSTREAM_HEAD_SHA",
"UpstreamCommitRecord",
"is_verified_check_conclusions",
"schedule_upstream_check",
"select_latest_verified_commit",
]
corridor_key/config.py
+51
View File
@@ -0,0 +1,51 @@
from __future__ import annotations
from dataclasses import dataclass
@dataclass(frozen=True, slots=True)
class CorridorKeySettings:
gamma_space: str = "sRGB"
feather_radius: int = 4
edge_focus: float = 1.5
threshold: float = 0.5
preserve_core: float = 0.2
despill_strength: float = 1.0
refiner_strength: float = 1.0
auto_despeckle: str = "On"
despeckle_size: int = 400
output_mode: str = "Processed"
def __post_init__(self) -> None:
if self.gamma_space not in {"sRGB", "Linear"}:
raise ValueError("gamma_space must be 'sRGB' or 'Linear'.")
if not 0 <= self.feather_radius <= 64:
raise ValueError("feather_radius must be between 0 and 64.")
if not 0.0 <= self.edge_focus <= 4.0:
raise ValueError("edge_focus must be between 0.0 and 4.0.")
if not 0.0 <= self.threshold <= 1.0:
raise ValueError("threshold must be between 0.0 and 1.0.")
if not 0.0 <= self.preserve_core <= 0.45:
raise ValueError("preserve_core must be between 0.0 and 0.45.")
if not 0.0 <= self.despill_strength <= 1.0:
raise ValueError("despill_strength must be between 0.0 and 1.0.")
if not 0.1 <= self.refiner_strength <= 4.0:
raise ValueError("refiner_strength must be between 0.1 and 4.0.")
if self.auto_despeckle not in {"Off", "On"}:
raise ValueError("auto_despeckle must be 'Off' or 'On'.")
if not 1 <= self.despeckle_size <= 4096:
raise ValueError("despeckle_size must be between 1 and 4096.")
if self.output_mode not in {"FG", "Matte", "Processed", "Comp"}:
raise ValueError("output_mode must be one of: FG, Matte, Processed, Comp.")
@property
def transition_gain(self) -> float:
return 8.0 + (self.edge_focus * 16.0)
@property
def despill_factor(self) -> float:
return self.despill_strength
@property
def despeckle_enabled(self) -> bool:
return self.auto_despeckle == "On"
corridor_key/processor.py
+133
View File
@@ -0,0 +1,133 @@
from __future__ import annotations
import logging
import torch
from .config import CorridorKeySettings
from .tensor_ops import (
build_checkerboard,
clamp_unit_interval,
compute_edge_map,
convert_gamma_space,
despeckle_binary_mask,
ensure_image_tensor,
ensure_mask_batch,
gaussian_blur_map,
mask_to_image,
)
LOGGER = logging.getLogger(__name__)
class CorridorKeyProcessor:
# Mini title: lightweight green-spill suppression for edge pixels.
def _apply_despill(self, image: torch.Tensor, alpha: torch.Tensor, strength: float) -> torch.Tensor:
if strength <= 0.0:
return image
edge_weight = 1.0 - (alpha.mul(2.0).sub(1.0).abs())
rb_max = torch.maximum(image[..., 0], image[..., 2])
spill = torch.relu(image[..., 1] - rb_max)
reduction = spill * strength * (0.25 + (0.75 * edge_weight))
red = clamp_unit_interval(image[..., 0] + (reduction * 0.10))
green = clamp_unit_interval(image[..., 1] - reduction)
blue = clamp_unit_interval(image[..., 2] + (reduction * 0.10))
return torch.stack((red, green, blue), dim=-1)
# Mini title: choose the user-requested primary preview.
def _select_output(
self,
output_mode: str,
alpha: torch.Tensor,
foreground: torch.Tensor,
processed: torch.Tensor,
comp: torch.Tensor,
) -> torch.Tensor:
if output_mode == "FG":
return foreground
if output_mode == "Matte":
return mask_to_image(alpha)
if output_mode == "Comp":
return comp
return processed
def refine(
self,
image: torch.Tensor,
mask: torch.Tensor,
settings: CorridorKeySettings,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
if not isinstance(settings, CorridorKeySettings):
raise ValueError("settings must be a CorridorKeySettings instance.")
image_batch = ensure_image_tensor(image)
working_image = convert_gamma_space(image_batch, settings.gamma_space)
mask_batch = ensure_mask_batch(
mask=mask,
batch_size=working_image.shape[0],
height=working_image.shape[1],
width=working_image.shape[2],
)
# Mini title: compute the soft starting matte.
softened_mask = gaussian_blur_map(mask_batch, settings.feather_radius)
# Mini title: derive edge guidance from image luminance.
edge_map = compute_edge_map(working_image)
edge_gain = 1.0 + (edge_map * settings.edge_focus)
centered_mask = (
(softened_mask - settings.threshold)
* settings.transition_gain
* settings.refiner_strength
)
alpha = torch.sigmoid(centered_mask * edge_gain)
low_lock = max(0.0, settings.threshold - settings.preserve_core)
high_lock = min(1.0, settings.threshold + settings.preserve_core)
alpha = torch.where(mask_batch <= low_lock, torch.zeros_like(alpha), alpha)
alpha = torch.where(mask_batch >= high_lock, torch.ones_like(alpha), alpha)
# Mini title: suppress tiny disconnected islands when requested.
if settings.despeckle_enabled:
support_mask = despeckle_binary_mask(
(alpha >= settings.threshold).to(dtype=alpha.dtype),
settings.despeckle_size,
)
alpha = torch.where(support_mask > 0.0, alpha, torch.zeros_like(alpha))
alpha = clamp_unit_interval(alpha)
# Mini title: build the foreground-style outputs.
despilled_image = self._apply_despill(working_image, alpha, settings.despill_factor)
visible_support = (alpha > 0.01).to(dtype=despilled_image.dtype).unsqueeze(-1)
foreground = clamp_unit_interval(despilled_image * visible_support)
alpha_rgb = alpha.unsqueeze(-1)
processed = clamp_unit_interval(foreground * alpha_rgb)
checkerboard = build_checkerboard(
batch_size=working_image.shape[0],
height=working_image.shape[1],
width=working_image.shape[2],
device=working_image.device,
dtype=working_image.dtype,
)
comp = clamp_unit_interval(processed + (checkerboard * (1.0 - alpha_rgb)))
selected_output = self._select_output(
output_mode=settings.output_mode,
alpha=alpha,
foreground=foreground,
processed=processed,
comp=comp,
)
LOGGER.debug(
"CorridorKey refinement complete for batch=%s height=%s width=%s",
working_image.shape[0],
working_image.shape[1],
working_image.shape[2],
)
return alpha, selected_output, foreground, processed, comp
corridor_key/tensor_ops.py
+177
View File
@@ -0,0 +1,177 @@
from __future__ import annotations
from collections import deque
import numpy as np
import torch
import torch.nn.functional as F
def clamp_unit_interval(tensor: torch.Tensor) -> torch.Tensor:
return tensor.to(dtype=torch.float32).clamp(0.0, 1.0)
def convert_gamma_space(image: torch.Tensor, gamma_space: str) -> torch.Tensor:
normalized = clamp_unit_interval(image)
if gamma_space == "Linear":
return normalized.clamp_min(0.0).pow(1.0 / 2.2)
return normalized
def ensure_image_tensor(image: torch.Tensor) -> torch.Tensor:
if not isinstance(image, torch.Tensor):
raise ValueError("image must be a torch.Tensor.")
if image.ndim != 4 or image.shape[-1] != 3:
raise ValueError("image must have shape [batch, height, width, 3].")
if image.shape[0] < 1 or image.shape[1] < 1 or image.shape[2] < 1:
raise ValueError("image must contain at least one pixel.")
return clamp_unit_interval(image)
def ensure_mask_batch(
mask: torch.Tensor,
batch_size: int,
height: int,
width: int,
) -> torch.Tensor:
if not isinstance(mask, torch.Tensor):
raise ValueError("mask must be a torch.Tensor.")
normalized = mask.to(dtype=torch.float32)
if normalized.ndim == 2:
normalized = normalized.unsqueeze(0)
elif normalized.ndim == 3:
pass
elif normalized.ndim == 4 and normalized.shape[1] == 1:
normalized = normalized.squeeze(1)
elif normalized.ndim == 4 and normalized.shape[-1] == 1:
normalized = normalized.squeeze(-1)
else:
raise ValueError("mask must have shape [height, width], [batch, height, width], or a single-channel 4D shape.")
if normalized.ndim != 3:
raise ValueError("mask normalization failed to produce a [batch, height, width] tensor.")
if normalized.shape[-2:] != (height, width):
raise ValueError("mask height and width must match the image tensor.")
if normalized.shape[0] == 1 and batch_size > 1:
normalized = normalized.repeat(batch_size, 1, 1)
elif normalized.shape[0] != batch_size:
raise ValueError("mask batch dimension must match the image batch or be a single mask.")
return clamp_unit_interval(normalized)
def gaussian_blur_map(mask_batch: torch.Tensor, radius: int) -> torch.Tensor:
if radius <= 0:
return clamp_unit_interval(mask_batch)
sigma = max(radius / 2.0, 0.5)
positions = torch.arange(-radius, radius + 1, device=mask_batch.device, dtype=mask_batch.dtype)
kernel_1d = torch.exp(-(positions**2) / (2 * sigma * sigma))
kernel_1d = kernel_1d / kernel_1d.sum()
kernel_x = kernel_1d.view(1, 1, 1, -1)
kernel_y = kernel_1d.view(1, 1, -1, 1)
working = mask_batch.unsqueeze(1)
working = F.pad(working, (radius, radius, 0, 0), mode="replicate")
working = F.conv2d(working, kernel_x)
working = F.pad(working, (0, 0, radius, radius), mode="replicate")
working = F.conv2d(working, kernel_y)
return clamp_unit_interval(working.squeeze(1))
def despeckle_binary_mask(mask_batch: torch.Tensor, minimum_size: int) -> torch.Tensor:
if minimum_size <= 1:
return clamp_unit_interval(mask_batch)
binary = (clamp_unit_interval(mask_batch) >= 0.5).to(dtype=torch.uint8)
output = binary.clone()
cpu_masks = binary.detach().cpu().numpy()
for batch_index, cpu_mask in enumerate(cpu_masks):
keep_mask = np.zeros_like(cpu_mask, dtype=np.uint8)
visited = np.zeros_like(cpu_mask, dtype=bool)
height, width = cpu_mask.shape
for y_coord in range(height):
for x_coord in range(width):
if cpu_mask[y_coord, x_coord] == 0 or visited[y_coord, x_coord]:
continue
queue: deque[tuple[int, int]] = deque([(y_coord, x_coord)])
component: list[tuple[int, int]] = []
visited[y_coord, x_coord] = True
while queue:
current_y, current_x = queue.popleft()
component.append((current_y, current_x))
for next_y, next_x in (
(current_y - 1, current_x),
(current_y + 1, current_x),
(current_y, current_x - 1),
(current_y, current_x + 1),
):
if next_y < 0 or next_y >= height or next_x < 0 or next_x >= width:
continue
if visited[next_y, next_x] or cpu_mask[next_y, next_x] == 0:
continue
visited[next_y, next_x] = True
queue.append((next_y, next_x))
if len(component) >= minimum_size:
for keep_y, keep_x in component:
keep_mask[keep_y, keep_x] = 1
output[batch_index] = torch.from_numpy(keep_mask).to(device=mask_batch.device)
return output.to(dtype=mask_batch.dtype)
def compute_edge_map(image: torch.Tensor) -> torch.Tensor:
luminance = (
(image[..., 0] * 0.299)
+ (image[..., 1] * 0.587)
+ (image[..., 2] * 0.114)
).unsqueeze(1)
sobel_x = image.new_tensor(
[[[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]]]
).unsqueeze(0)
sobel_y = image.new_tensor(
[[[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0]]]
).unsqueeze(0)
padded = F.pad(luminance, (1, 1, 1, 1), mode="replicate")
grad_x = F.conv2d(padded, sobel_x)
grad_y = F.conv2d(padded, sobel_y)
magnitude = torch.sqrt((grad_x.square() + grad_y.square()).clamp_min(1e-12)).squeeze(1)
batch_max = magnitude.flatten(1).max(dim=1).values.clamp_min(1e-6).view(-1, 1, 1)
return clamp_unit_interval(magnitude / batch_max)
def mask_to_image(mask_batch: torch.Tensor) -> torch.Tensor:
grayscale = clamp_unit_interval(mask_batch).unsqueeze(-1)
return grayscale.repeat(1, 1, 1, 3)
def build_checkerboard(
batch_size: int,
height: int,
width: int,
device: torch.device,
dtype: torch.dtype,
tile_size: int = 8,
) -> torch.Tensor:
y_index = torch.arange(height, device=device)
x_index = torch.arange(width, device=device)
grid_y = (y_index[:, None] // tile_size) % 2
grid_x = (x_index[None, :] // tile_size) % 2
board = ((grid_x + grid_y) % 2).to(dtype=dtype)
board = (board * 0.30) + 0.35
board = board.unsqueeze(0).unsqueeze(-1).repeat(batch_size, 1, 1, 3)
return clamp_unit_interval(board)
corridor_key/upstream_sync.py
+197
View File
@@ -0,0 +1,197 @@
from __future__ import annotations
import json
import logging
import os
import threading
from dataclasses import dataclass
from typing import Any, Sequence
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen
LOGGER = logging.getLogger(__name__)
UPSTREAM_API_ROOT = "https://api.github.com/repos/nikopueringer/CorridorKey"
SYNCED_UPSTREAM_HEAD_SHA = "38248c6dd1785b8d7e51ef257fd0889143ca58af"
SYNCED_UPSTREAM_HEAD_DATE = "2026-03-01T22:23:57Z"
SYNCED_UPSTREAM_HEAD_MESSAGE = (
"Merge pull request #34 from taylorOntologize/tgw/coverage-review"
)
SYNCED_UPSTREAM_HEAD_CHECK_CONCLUSIONS = ("failure", "success", "cancelled")
_BLOCKING_CHECK_CONCLUSIONS = {
"action_required",
"failure",
"startup_failure",
"timed_out",
}
_CHECK_THREAD_STARTED = False
_CHECK_THREAD_LOCK = threading.Lock()
@dataclass(frozen=True, slots=True)
class UpstreamCommitRecord:
sha: str
date: str
message: str
conclusions: tuple[str, ...]
def _parse_bool_env(name: str, default: bool) -> bool:
raw_value = os.getenv(name)
if raw_value is None:
return default
return raw_value.strip().lower() not in {"0", "false", "no", "off"}
def _parse_float_env(name: str, default: float, minimum: float, maximum: float) -> float:
raw_value = os.getenv(name)
if raw_value is None:
return default
try:
parsed = float(raw_value)
except ValueError:
return default
return max(minimum, min(maximum, parsed))
def _parse_int_env(name: str, default: int, minimum: int, maximum: int) -> int:
raw_value = os.getenv(name)
if raw_value is None:
return default
try:
parsed = int(raw_value)
except ValueError:
return default
return max(minimum, min(maximum, parsed))
def is_verified_check_conclusions(conclusions: Sequence[str]) -> bool:
normalized = tuple(
str(conclusion).strip().lower()
for conclusion in conclusions
if str(conclusion).strip()
)
if not normalized:
return False
has_success = any(conclusion == "success" for conclusion in normalized)
has_blocker = any(conclusion in _BLOCKING_CHECK_CONCLUSIONS for conclusion in normalized)
return has_success and not has_blocker
def _extract_commit_record(commit_payload: dict[str, Any], check_payload: dict[str, Any]) -> UpstreamCommitRecord:
message = str(
commit_payload.get("commit", {})
.get("message", "")
).splitlines()[0]
conclusions = tuple(
str(check_run.get("conclusion", "")).strip().lower()
for check_run in check_payload.get("check_runs", [])
)
return UpstreamCommitRecord(
sha=str(commit_payload.get("sha", "")).strip(),
date=str(commit_payload.get("commit", {}).get("author", {}).get("date", "")).strip(),
message=message,
conclusions=conclusions,
)
def select_latest_verified_commit(
commit_payloads: Sequence[dict[str, Any]],
check_payloads_by_sha: dict[str, dict[str, Any]],
) -> UpstreamCommitRecord | None:
for commit_payload in commit_payloads:
sha = str(commit_payload.get("sha", "")).strip()
if not sha:
continue
check_payload = check_payloads_by_sha.get(sha)
if check_payload is None:
continue
record = _extract_commit_record(commit_payload, check_payload)
if is_verified_check_conclusions(record.conclusions):
return record
return None
def _fetch_json(url: str, timeout_seconds: float) -> Any:
request = Request(
url,
headers={
"Accept": "application/vnd.github+json",
"User-Agent": "ComfyUI-CorridorKey",
},
method="GET",
)
with urlopen(request, timeout=timeout_seconds) as response:
return json.loads(response.read().decode("utf-8"))
def fetch_latest_verified_commit(timeout_seconds: float, depth: int) -> UpstreamCommitRecord | None:
commits_url = f"{UPSTREAM_API_ROOT}/commits?sha=main&per_page={depth}"
commit_payloads = _fetch_json(commits_url, timeout_seconds)
if not isinstance(commit_payloads, list):
raise ValueError("GitHub commits API returned an unexpected payload.")
check_payloads_by_sha: dict[str, dict[str, Any]] = {}
for commit_payload in commit_payloads:
sha = str(commit_payload.get("sha", "")).strip()
if not sha:
continue
check_url = f"{UPSTREAM_API_ROOT}/commits/{sha}/check-runs"
check_payload = _fetch_json(check_url, timeout_seconds)
if isinstance(check_payload, dict):
check_payloads_by_sha[sha] = check_payload
return select_latest_verified_commit(commit_payloads, check_payloads_by_sha)
def _run_upstream_check() -> None:
timeout_seconds = _parse_float_env("CORRIDORKEY_UPSTREAM_TIMEOUT_SECONDS", 3.0, 0.5, 10.0)
depth = _parse_int_env("CORRIDORKEY_UPSTREAM_CHECK_DEPTH", 15, 1, 30)
try:
latest_verified = fetch_latest_verified_commit(
timeout_seconds=timeout_seconds,
depth=depth,
)
except (HTTPError, URLError, OSError, TimeoutError, ValueError) as exc:
LOGGER.debug("CorridorKey upstream check skipped: %s", exc)
return
if latest_verified is None:
LOGGER.debug(
"CorridorKey upstream check found no recent commit without failing checks."
)
return
if latest_verified.sha == SYNCED_UPSTREAM_HEAD_SHA:
LOGGER.debug(
"CorridorKey upstream head %s is already the current reviewed baseline.",
latest_verified.sha[:12],
)
return
LOGGER.info(
"CorridorKey upstream update available: %s (%s). "
"This node does not auto-overwrite itself; review and port changes manually.",
latest_verified.sha[:12],
latest_verified.date,
)
def schedule_upstream_check() -> None:
if not _parse_bool_env("CORRIDORKEY_AUTO_CHECK_UPSTREAM", True):
return
global _CHECK_THREAD_STARTED
with _CHECK_THREAD_LOCK:
if _CHECK_THREAD_STARTED:
return
_CHECK_THREAD_STARTED = True
thread = threading.Thread(
target=_run_upstream_check,
name="CorridorKeyUpstreamCheck",
daemon=True,
)
thread.start()
nodes.py
+120
View File
@@ -0,0 +1,120 @@
from __future__ import annotations
try:
from .corridor_key import CorridorKeyProcessor, CorridorKeySettings
except ImportError:
from corridor_key import CorridorKeyProcessor, CorridorKeySettings
class CorridorKey:
@classmethod
def INPUT_TYPES(cls) -> dict[str, dict[str, tuple]]:
return {
"required": {
"image": ("IMAGE",),
"mask": ("MASK",),
"gamma_space": (["sRGB", "Linear"],),
"feather_radius": (
"INT",
{
"default": 4,
"min": 0,
"max": 64,
"step": 1,
},
),
"edge_focus": (
"FLOAT",
{
"default": 1.5,
"min": 0.0,
"max": 4.0,
"step": 0.1,
},
),
"threshold": (
"FLOAT",
{
"default": 0.5,
"min": 0.0,
"max": 1.0,
"step": 0.01,
},
),
"preserve_core": (
"FLOAT",
{
"default": 0.2,
"min": 0.0,
"max": 0.45,
"step": 0.01,
},
),
"despill_strength": (
"FLOAT",
{
"default": 1.0,
"min": 0.0,
"max": 1.0,
"step": 0.01,
},
),
"refiner_strength": (
"FLOAT",
{
"default": 1.0,
"min": 0.1,
"max": 4.0,
"step": 0.1,
},
),
"auto_despeckle": (["On", "Off"],),
"despeckle_size": (
"INT",
{
"default": 400,
"min": 1,
"max": 4096,
"step": 1,
},
),
"output_mode": (["Processed", "FG", "Matte", "Comp"],),
}
}
RETURN_TYPES = ("MASK", "IMAGE", "IMAGE", "IMAGE", "IMAGE")
RETURN_NAMES = ("refined_mask", "selected_output", "foreground", "processed", "comp")
FUNCTION = "run"
CATEGORY = "CorridorKey"
def __init__(self) -> None:
self._processor = CorridorKeyProcessor()
def run(
self,
image,
mask,
gamma_space: str,
feather_radius: int,
edge_focus: float,
threshold: float,
preserve_core: float,
despill_strength: float,
refiner_strength: float,
auto_despeckle: str,
despeckle_size: int,
output_mode: str,
):
settings = CorridorKeySettings(
gamma_space=str(gamma_space),
feather_radius=int(feather_radius),
edge_focus=float(edge_focus),
threshold=float(threshold),
preserve_core=float(preserve_core),
despill_strength=float(despill_strength),
refiner_strength=float(refiner_strength),
auto_despeckle=str(auto_despeckle),
despeckle_size=int(despeckle_size),
output_mode=str(output_mode),
)
return self._processor.refine(image=image, mask=mask, settings=settings)
pyproject.toml
+46
View File
@@ -0,0 +1,46 @@
[build-system]
requires = ["setuptools>=68", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "comfyui-corridorkey"
version = "0.1.0"
description = "ComfyUI custom node for CorridorKey-style edge-aware coarse mask refinement."
readme = "README.md"
requires-python = ">=3.10"
dependencies = [
# Keep runtime deps aligned with requirements.txt for ComfyUI installs.
"numpy>=1.24",
"Pillow>=9.5",
"torch>=2.0",
]
[project.optional-dependencies]
dev = [
"black>=24.0.0",
"mypy>=1.10.0",
"pytest>=8.0.0",
"ruff>=0.6.0",
]
[tool.setuptools]
packages = ["corridor_key"]
[tool.black]
line-length = 100
target-version = ["py310"]
[tool.ruff]
line-length = 100
target-version = "py310"
[tool.ruff.lint]
select = ["E", "F", "I", "B"]
[tool.mypy]
python_version = "3.10"
strict = false
warn_unused_configs = true
disallow_untyped_defs = true
disallow_incomplete_defs = true
no_implicit_optional = true
requirements.txt
+7
View File
@@ -0,0 +1,7 @@
# Runtime dependencies for ComfyUI-CorridorKey
# ComfyUI usually already ships with a suitable torch build. Keeping it listed
# here matches common custom-node conventions and lets dependency installers
# verify it is present without changing the node code.
torch>=2.0
numpy>=1.24
Pillow>=9.5
tests/test_processor.py
+206
View File
@@ -0,0 +1,206 @@
from __future__ import annotations
import pytest
import torch
from corridor_key import (
CorridorKeyProcessor,
CorridorKeySettings,
is_verified_check_conclusions,
select_latest_verified_commit,
)
def make_image(batch: int = 1, height: int = 6, width: int = 6) -> torch.Tensor:
image = torch.zeros((batch, height, width, 3), dtype=torch.float32)
image[..., : width // 2, :] = 0.15
image[..., width // 2 :, :] = 0.85
return image
def test_settings_reject_invalid_threshold() -> None:
with pytest.raises(ValueError, match="threshold"):
CorridorKeySettings(threshold=1.2)
def test_settings_match_upstream_like_defaults() -> None:
settings = CorridorKeySettings()
assert settings.gamma_space == "sRGB"
assert settings.despill_strength == 1.0
assert settings.refiner_strength == 1.0
assert settings.auto_despeckle == "On"
assert settings.despeckle_size == 400
assert settings.output_mode == "Processed"
def test_refine_broadcasts_single_mask_across_batch() -> None:
processor = CorridorKeyProcessor()
image = make_image(batch=2)
mask = torch.zeros((6, 6), dtype=torch.float32)
mask[:, 3:] = 1.0
refined_mask, selected_output, foreground, processed, comp = processor.refine(
image=image,
mask=mask,
settings=CorridorKeySettings(auto_despeckle="Off"),
)
assert refined_mask.shape == (2, 6, 6)
assert selected_output.shape == image.shape
assert foreground.shape == image.shape
assert processed.shape == image.shape
assert comp.shape == image.shape
assert torch.allclose(refined_mask[0], refined_mask[1])
def test_refine_preserves_full_foreground() -> None:
processor = CorridorKeyProcessor()
image = make_image()
mask = torch.ones((1, 6, 6), dtype=torch.float32)
refined_mask, selected_output, foreground, processed, comp = processor.refine(
image=image,
mask=mask,
settings=CorridorKeySettings(auto_despeckle="Off"),
)
assert torch.allclose(refined_mask, torch.ones_like(refined_mask))
assert torch.allclose(foreground, image)
assert torch.allclose(processed, image)
assert torch.allclose(selected_output, processed)
assert torch.allclose(comp, image)
def test_refine_preserves_full_background() -> None:
processor = CorridorKeyProcessor()
image = make_image()
mask = torch.zeros((1, 6, 6), dtype=torch.float32)
refined_mask, selected_output, foreground, processed, comp = processor.refine(
image=image,
mask=mask,
settings=CorridorKeySettings(auto_despeckle="Off"),
)
assert torch.allclose(refined_mask, torch.zeros_like(refined_mask))
assert torch.allclose(foreground, torch.zeros_like(foreground))
assert torch.allclose(processed, torch.zeros_like(processed))
assert not torch.allclose(comp, torch.zeros_like(comp))
assert torch.allclose(selected_output, processed)
def test_refine_rejects_wrong_mask_size() -> None:
processor = CorridorKeyProcessor()
image = make_image(height=6, width=6)
mask = torch.zeros((1, 5, 6), dtype=torch.float32)
with pytest.raises(ValueError, match="height and width"):
processor.refine(
image=image,
mask=mask,
settings=CorridorKeySettings(auto_despeckle="Off"),
)
def test_output_mode_matte_returns_grayscale_preview() -> None:
processor = CorridorKeyProcessor()
image = make_image()
mask = torch.zeros((1, 6, 6), dtype=torch.float32)
mask[:, :, 3:] = 1.0
refined_mask, selected_output, _, _, _ = processor.refine(
image=image,
mask=mask,
settings=CorridorKeySettings(output_mode="Matte", auto_despeckle="Off"),
)
assert torch.allclose(selected_output[..., 0], refined_mask)
assert torch.allclose(selected_output[..., 1], refined_mask)
assert torch.allclose(selected_output[..., 2], refined_mask)
def test_despill_reduces_green_bias() -> None:
processor = CorridorKeyProcessor()
image = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
image[..., 0] = 0.1
image[..., 1] = 0.9
image[..., 2] = 0.1
mask = torch.ones((1, 4, 4), dtype=torch.float32)
_, _, foreground_without, _, _ = processor.refine(
image=image,
mask=mask,
settings=CorridorKeySettings(despill_strength=0.0, auto_despeckle="Off"),
)
_, _, foreground_with, _, _ = processor.refine(
image=image,
mask=mask,
settings=CorridorKeySettings(despill_strength=1.0, auto_despeckle="Off"),
)
assert torch.mean(foreground_with[..., 1]) < torch.mean(foreground_without[..., 1])
def test_auto_despeckle_can_remove_single_pixel_island() -> None:
processor = CorridorKeyProcessor()
image = make_image()
mask = torch.zeros((1, 6, 6), dtype=torch.float32)
mask[:, 2, 2] = 1.0
refined_mask, _, _, _, _ = processor.refine(
image=image,
mask=mask,
settings=CorridorKeySettings(
auto_despeckle="On",
despeckle_size=9,
feather_radius=0,
preserve_core=0.0,
),
)
assert float(refined_mask[0, 2, 2]) == 0.0
def test_verified_check_policy_requires_success_and_no_failure() -> None:
assert is_verified_check_conclusions(("success", "cancelled"))
assert not is_verified_check_conclusions(("success", "failure"))
assert not is_verified_check_conclusions(("cancelled",))
def test_select_latest_verified_commit_picks_first_clean_candidate() -> None:
commit_payloads = [
{
"sha": "newest-bad",
"commit": {
"author": {"date": "2026-03-02T00:00:00Z"},
"message": "Newest bad commit",
},
},
{
"sha": "next-good",
"commit": {
"author": {"date": "2026-03-01T00:00:00Z"},
"message": "Next good commit",
},
},
]
check_payloads_by_sha = {
"newest-bad": {
"check_runs": [
{"conclusion": "success"},
{"conclusion": "failure"},
]
},
"next-good": {
"check_runs": [
{"conclusion": "success"},
{"conclusion": "cancelled"},
]
},
}
selected = select_latest_verified_commit(commit_payloads, check_payloads_by_sha)
assert selected is not None
assert selected.sha == "next-good"