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

Initial commit: ComfyUI nodes for CorridorKey greenscreen keyer

Browse Source 
Wraps Niko Pueringer's CorridorKey neural network keyer as ComfyUI
custom nodes. Includes Load Model, Greenscreen, Chroma Key Prepass,
and Composite nodes. CorridorKey is installed as a pip dependency
with auto-download of the model checkpoint from HuggingFace.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pixelworld AI
2026-03-18 08:25:41 -05:00
co-authored by Claude Opus 4.6
commit 25ad75758e
9 changed files with 543 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+3
View File
@@ -0,0 +1,3 @@
__pycache__/
*.pyc
CorridorKey/
README.md
+50
View File
@@ -0,0 +1,50 @@
# ComfyUI-CorridorKey
Custom ComfyUI nodes for [CorridorKey](https://github.com/nikopueringer/CorridorKey) — the neural network green screen keyer by [Niko Pueringer](https://github.com/nikopueringer) and [Corridor Crew](https://www.youtube.com/@Corridor).
CorridorKey solves the color unmixing problem that traditional keyers struggle with. For every pixel — even highly transparent ones like motion blur or out-of-focus edges — the model predicts the true foreground color and a clean alpha channel, as if the green screen was never there.
## Nodes
| Node | Description |
|------|-------------|
| **Load CorridorKey Model** | Load the GreenFormer neural network. The model checkpoint (~300MB) is auto-downloaded from HuggingFace on first use. |
| **CorridorKey Greenscreen** | Run inference on green screen footage. Takes an image and a coarse alpha hint, outputs a clean alpha, foreground, RGBA, and composite preview. |
| **Chroma Key Prepass** | Fast chroma key that generates a coarse alpha hint from green/blue screen footage. Designed as input to CorridorKey Greenscreen. |
| **CorridorKey Composite** | Composite a keyed foreground over a background with proper linear-space blending and sRGB output. |
## Installation
### Via ComfyUI Manager
Search for **ComfyUI-CorridorKey** in the ComfyUI Manager and install.
### Manual
```bash
cd ComfyUI/custom_nodes
git clone https://github.com/pixelworldai/ComfyUI-CorridorKey.git
cd ComfyUI-CorridorKey
pip install -r requirements.txt
```
The CorridorKey model checkpoint is automatically downloaded to `ComfyUI/models/corridorkey/` on first use. You can also place a `.pth` file there manually.
## Usage
1. **Load CorridorKey Model** — loads the model onto your GPU
2. **Chroma Key Prepass** — feed in your green screen footage to generate a rough alpha hint (or bring your own from any roto/keying tool)
3. **CorridorKey Greenscreen** — connect the model, your image, and the alpha hint to get a clean key
4. **CorridorKey Composite** (optional) — composite the result over a background
The alpha hint does not need to be precise. CorridorKey was trained on coarse, blurry, eroded masks and excels at filling in fine detail from rough hints.
## Credits & Acknowledgements
**CorridorKey** was created by [Niko Pueringer](https://github.com/nikopueringer) at [Corridor Digital](https://corridordigital.com). The core model, inference engine, and all neural network architecture are entirely his work. This node pack is a ComfyUI wrapper around the original [CorridorKey](https://github.com/nikopueringer/CorridorKey) project.
Join the [Corridor Creates Discord](https://discord.gg/zvwUrdWXJm) to share ideas, results, and contribute to CorridorKey's development.
## License
The CorridorKey model and engine are licensed under the [Corridor Key Licence](https://github.com/nikopueringer/CorridorKey/blob/main/LICENSE) (CC BY-NC-SA 4.0 with additional terms). You may use this tool for processing images as part of commercial projects, but you may not repackage and sell the tool itself or offer inference as a paid API service. See the full license for details.
__init__.py
+10
View File
@@ -0,0 +1,10 @@
"""ComfyUI-CorridorKey: Neural network greenscreen keying nodes for ComfyUI."""
import os
# OpenEXR support must be enabled before any OpenCV import
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
nodes/__init__.py
+18
View File
@@ -0,0 +1,18 @@
from .load_model import LoadCorridorKeyModel
from .greenscreen import CorridorKeyGreenscreen
from .composite import CorridorKeyComposite
from .chroma_prepass import ChromaKeyPrepass
NODE_CLASS_MAPPINGS = {
"LoadCorridorKeyModel": LoadCorridorKeyModel,
"CorridorKeyGreenscreen": CorridorKeyGreenscreen,
"CorridorKeyComposite": CorridorKeyComposite,
"ChromaKeyPrepass": ChromaKeyPrepass,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"LoadCorridorKeyModel": "Load CorridorKey Model",
"CorridorKeyGreenscreen": "CorridorKey Greenscreen",
"CorridorKeyComposite": "CorridorKey Composite",
"ChromaKeyPrepass": "Chroma Key Prepass",
}
nodes/chroma_prepass.py
+116
View File
@@ -0,0 +1,116 @@
from __future__ import annotations
import logging
import time
import numpy as np
import torch
from comfy.utils import ProgressBar
logger = logging.getLogger("CorridorKey")
class ChromaKeyPrepass:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
},
"optional": {
"screen_color": (["green", "blue"], {"default": "green"}),
"threshold": ("FLOAT", {
"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "How much the screen channel must exceed the others to be considered background. Lower = more aggressive keying.",
}),
"softness": ("FLOAT", {
"default": 0.1, "min": 0.0, "max": 0.5, "step": 0.01,
"tooltip": "Width of the soft falloff between foreground and background. 0 = hard edge.",
}),
"erode": ("INT", {
"default": 2, "min": 0, "max": 50, "step": 1,
"tooltip": "Shrink the foreground mask by this many pixels. Helps CorridorKey by pulling the hint away from edges.",
}),
"blur": ("INT", {
"default": 5, "min": 0, "max": 50, "step": 1,
"tooltip": "Gaussian blur radius on the final mask. Makes the hint coarser, which CorridorKey handles well.",
}),
},
}
RETURN_TYPES = ("MASK",)
RETURN_NAMES = ("alpha_hint",)
FUNCTION = "generate"
CATEGORY = "CorridorKey"
DESCRIPTION = (
"Fast chroma key pre-pass that generates a coarse alpha hint from "
"green/blue screen footage. Designed as input to CorridorKey Greenscreen — "
"does not need to be precise."
)
def generate(
self,
image: torch.Tensor,
screen_color: str = "green",
threshold: float = 0.3,
softness: float = 0.1,
erode: int = 2,
blur: int = 5,
):
import cv2
batch_size = image.shape[0]
logger.info(
"ChromaPrepass: %d frame%s | screen=%s threshold=%.2f softness=%.2f erode=%d blur=%d",
batch_size, "s" if batch_size > 1 else "",
screen_color, threshold, softness, erode, blur,
)
image = image.float()
pbar = ProgressBar(batch_size)
start_time = time.perf_counter()
masks = []
for i in range(batch_size):
img = image[i].cpu().numpy() # [H, W, 3] float32 0-1
r, g, b = img[:, :, 0], img[:, :, 1], img[:, :, 2]
# Difference key: how much does the screen channel exceed the others?
if screen_color == "green":
diff = g - np.maximum(r, b)
else:
diff = b - np.maximum(r, g)
# Soft threshold: background where diff > threshold
if softness > 0.001:
# Smooth ramp from 1 (foreground) to 0 (background)
alpha = 1.0 - np.clip((diff - threshold + softness) / (2.0 * softness), 0.0, 1.0)
else:
alpha = np.where(diff > threshold, 0.0, 1.0).astype(np.float32)
alpha = alpha.astype(np.float32)
# Erode to pull mask away from edges
if erode > 0:
kernel_size = erode * 2 + 1
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
alpha = cv2.erode(alpha, kernel)
# Blur for coarseness
if blur > 0:
blur_size = blur * 2 + 1
alpha = cv2.GaussianBlur(alpha, (blur_size, blur_size), 0)
alpha = np.clip(alpha, 0.0, 1.0)
masks.append(torch.from_numpy(alpha))
pbar.update(1)
elapsed = time.perf_counter() - start_time
logger.info(
"ChromaPrepass: finished %d frame%s in %.2fs",
batch_size, "s" if batch_size > 1 else "", elapsed,
)
mask_batch = torch.stack(masks, dim=0).float() # [B, H, W]
return (mask_batch,)
nodes/composite.py
+72
View File
@@ -0,0 +1,72 @@
from __future__ import annotations
import numpy as np
import torch
class CorridorKeyComposite:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"foreground": ("IMAGE",),
"alpha": ("MASK",),
"background": ("IMAGE",),
},
"optional": {
"fg_is_linear": ("BOOLEAN", {
"default": False,
"tooltip": "Set True if the foreground is already in linear color space.",
}),
},
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("composite",)
FUNCTION = "composite"
CATEGORY = "CorridorKey"
DESCRIPTION = (
"Composite a keyed foreground over a background with proper "
"linear-space blending and sRGB output."
)
def composite(
self,
foreground: torch.Tensor,
alpha: torch.Tensor,
background: torch.Tensor,
fg_is_linear: bool = False,
):
from CorridorKeyModule.core import color_utils as cu
batch_size = foreground.shape[0]
results = []
for i in range(batch_size):
fg_np = foreground[i].cpu().numpy().astype(np.float32)
bg_np = background[min(i, background.shape[0] - 1)].cpu().numpy().astype(np.float32)
# Alpha: [B,H,W] -> [H,W,1]
a_np = alpha[min(i, alpha.shape[0] - 1)].cpu().numpy().astype(np.float32)
if a_np.ndim == 2:
a_np = a_np[:, :, np.newaxis]
# Resize background to match foreground if needed
if bg_np.shape[:2] != fg_np.shape[:2]:
import cv2
bg_np = cv2.resize(bg_np, (fg_np.shape[1], fg_np.shape[0]), interpolation=cv2.INTER_LANCZOS4)
# Convert everything to linear for compositing
fg_lin = fg_np if fg_is_linear else cu.srgb_to_linear(fg_np)
bg_lin = cu.srgb_to_linear(bg_np)
# Straight-over composite in linear space
comp_lin = cu.composite_straight(fg_lin, bg_lin, a_np)
# Back to sRGB for ComfyUI display
comp_srgb = cu.linear_to_srgb(comp_lin)
comp_srgb = np.clip(comp_srgb, 0.0, 1.0)
results.append(torch.from_numpy(comp_srgb))
return (torch.stack(results, dim=0).float(),)
nodes/greenscreen.py
+155
View File
@@ -0,0 +1,155 @@
from __future__ import annotations
import logging
import time
import numpy as np
import torch
from comfy.utils import ProgressBar
logger = logging.getLogger("CorridorKey")
class CorridorKeyGreenscreen:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"ck_model": ("CK_MODEL",),
"image": ("IMAGE",),
"alpha_hint": ("MASK",),
},
"optional": {
"despill_strength": ("FLOAT", {
"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1,
"tooltip": "Strength of green spill removal. 0 = no despill.",
}),
"refiner_scale": ("FLOAT", {
"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05,
"tooltip": "Multiplier for the CNN refiner. 1.0 = default, 0.0 = disable refiner.",
}),
"auto_despeckle": ("BOOLEAN", {
"default": True,
"tooltip": "Remove small disconnected alpha islands (tracking markers, noise).",
}),
"despeckle_size": ("INT", {
"default": 400, "min": 0, "max": 10000, "step": 50,
"tooltip": "Minimum pixel area to keep when despeckling.",
}),
"input_is_linear": ("BOOLEAN", {
"default": False,
"tooltip": "Set True if your input image is in linear color space (e.g. from EXR loader).",
}),
},
}
RETURN_TYPES = ("MASK", "IMAGE", "IMAGE", "IMAGE")
RETURN_NAMES = ("alpha", "foreground", "processed_rgba", "comp_preview")
FUNCTION = "key"
CATEGORY = "CorridorKey"
DESCRIPTION = (
"Neural network greenscreen keying. Produces a clean alpha matte, "
"despilled straight foreground, premultiplied linear RGBA, and a "
"checkerboard composite preview."
)
def key(
self,
ck_model,
image: torch.Tensor,
alpha_hint: torch.Tensor,
despill_strength: float = 1.0,
refiner_scale: float = 1.0,
auto_despeckle: bool = True,
despeckle_size: int = 400,
input_is_linear: bool = False,
):
batch_size = image.shape[0]
h, w = image.shape[1], image.shape[2]
logger.info(
"CorridorKey: keying %d frame%s at %dx%d | despill=%.1f refiner=%.2f despeckle=%s(%d)",
batch_size, "s" if batch_size > 1 else "", w, h,
despill_strength, refiner_scale, auto_despeckle, despeckle_size,
)
alphas = []
fgs = []
processed_list = []
comps = []
# Ensure float32 — ComfyUI may pass BFloat16 tensors which numpy can't handle
image = image.float()
alpha_hint = alpha_hint.float()
pbar = ProgressBar(batch_size)
batch_start = time.perf_counter()
for i in range(batch_size):
frame_start = time.perf_counter()
# Convert ComfyUI IMAGE [B,H,W,C] -> numpy [H,W,C] float32
img_np = image[i].cpu().numpy()
# Convert ComfyUI MASK [B,H,W] -> numpy [H,W] float32
if alpha_hint.dim() == 3:
mask_np = alpha_hint[min(i, alpha_hint.shape[0] - 1)].cpu().numpy()
else:
# Single mask without batch dim
mask_np = alpha_hint.cpu().numpy()
result = ck_model.process_frame(
image=img_np,
mask_linear=mask_np,
refiner_scale=refiner_scale,
input_is_linear=input_is_linear,
despill_strength=despill_strength,
auto_despeckle=auto_despeckle,
despeckle_size=despeckle_size,
)
# alpha: [H,W,1] linear float -> MASK [H,W]
alpha_out = result["alpha"]
if alpha_out.ndim == 3:
alpha_out = alpha_out[:, :, 0]
alphas.append(torch.from_numpy(alpha_out))
# fg: [H,W,3] sRGB straight float -> IMAGE [H,W,C]
fgs.append(torch.from_numpy(result["fg"]))
# processed: [H,W,4] linear premultiplied RGBA -> IMAGE [H,W,C]
processed_list.append(torch.from_numpy(result["processed"]))
# comp: [H,W,3] sRGB checkerboard composite -> IMAGE [H,W,C]
comps.append(torch.from_numpy(result["comp"]))
frame_time = time.perf_counter() - frame_start
elapsed = time.perf_counter() - batch_start
avg_per_frame = elapsed / (i + 1)
remaining = avg_per_frame * (batch_size - i - 1)
logger.info(
"CorridorKey: frame %d/%d done in %.2fs (avg %.2fs/frame, ~%.0fs remaining)",
i + 1, batch_size, frame_time, avg_per_frame, remaining,
)
pbar.update(1)
total_time = time.perf_counter() - batch_start
logger.info(
"CorridorKey: finished %d frame%s in %.1fs (%.2fs avg/frame)",
batch_size, "s" if batch_size > 1 else "", total_time, total_time / batch_size,
)
# Stack into batched tensors
alpha_batch = torch.stack(alphas, dim=0).float() # [B, H, W]
fg_batch = torch.stack(fgs, dim=0).float() # [B, H, W, 3]
processed_batch = torch.stack(processed_list, dim=0).float() # [B, H, W, 4]
comp_batch = torch.stack(comps, dim=0).float() # [B, H, W, 3]
# Clamp to valid range
alpha_batch = torch.clamp(alpha_batch, 0.0, 1.0)
fg_batch = torch.clamp(fg_batch, 0.0, 1.0)
processed_batch = torch.clamp(processed_batch, 0.0, 1.0)
comp_batch = torch.clamp(comp_batch, 0.0, 1.0)
return (alpha_batch, fg_batch, processed_batch, comp_batch)
nodes/load_model.py
+117
View File
@@ -0,0 +1,117 @@
from __future__ import annotations
import glob
import logging
import os
import torch
import comfy.model_management
import folder_paths
logger = logging.getLogger(__name__)
# Register a "corridorkey" model folder so users can place checkpoints in
# ComfyUI/models/corridorkey/ (created automatically if missing).
_CK_MODEL_DIR = os.path.join(folder_paths.models_dir, "corridorkey")
os.makedirs(_CK_MODEL_DIR, exist_ok=True)
folder_paths.folder_names_and_paths["corridorkey"] = (
[_CK_MODEL_DIR],
folder_paths.supported_pt_extensions,
)
_HF_REPO_ID = "nikopueringer/CorridorKey_v1.0"
_HF_FILENAME = "CorridorKey_v1.0.pth"
_LOCAL_FILENAME = "CorridorKey.pth"
def _download_checkpoint() -> str:
"""Download the CorridorKey checkpoint from HuggingFace Hub."""
from huggingface_hub import hf_hub_download
logger.info("Downloading CorridorKey model from %s ...", _HF_REPO_ID)
downloaded = hf_hub_download(
repo_id=_HF_REPO_ID,
filename=_HF_FILENAME,
local_dir=_CK_MODEL_DIR,
)
# Rename to the canonical local name
dest = os.path.join(_CK_MODEL_DIR, _LOCAL_FILENAME)
if downloaded != dest:
os.replace(downloaded, dest)
logger.info("CorridorKey model saved to %s", dest)
return dest
def _find_checkpoint() -> str:
"""Find or auto-download the CorridorKey checkpoint.
Lookup order:
1. ComfyUI/models/corridorkey/*.pth (recommended)
2. <pip package>/CorridorKeyModule/checkpoints/*.pth (fallback)
3. Auto-download from HuggingFace Hub
"""
# 1. ComfyUI models directory
candidates = glob.glob(os.path.join(_CK_MODEL_DIR, "*.pth"))
if candidates:
return candidates[0]
# 2. Inside the pip-installed package
try:
import CorridorKeyModule
pkg_ckpt_dir = os.path.join(os.path.dirname(CorridorKeyModule.__file__), "checkpoints")
candidates = glob.glob(os.path.join(pkg_ckpt_dir, "*.pth"))
if candidates:
return candidates[0]
except ImportError:
pass
# 3. Auto-download
return _download_checkpoint()
class LoadCorridorKeyModel:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {},
"optional": {
"precision": (["auto", "fp16", "fp32"], {"default": "auto"}),
"use_refiner": ("BOOLEAN", {"default": True}),
"mixed_precision": ("BOOLEAN", {"default": True}),
},
}
RETURN_TYPES = ("CK_MODEL",)
RETURN_NAMES = ("ck_model",)
FUNCTION = "load_model"
CATEGORY = "CorridorKey"
DESCRIPTION = "Load the CorridorKey GreenFormer neural network model for greenscreen keying."
def load_model(self, precision="auto", use_refiner=True, mixed_precision=True):
from CorridorKeyModule import CorridorKeyEngine
checkpoint_path = _find_checkpoint()
device = comfy.model_management.get_torch_device()
# Model weights are always stored as fp32 — the engine's mixed_precision
# autocast handles fp16 inference automatically for speed.
# BFloat16 weights cause issues with numpy conversion downstream.
if precision == "fp16":
model_precision = torch.float16
else:
model_precision = torch.float32
logger.info("Loading CorridorKey on %s with %s precision", device, model_precision)
engine = CorridorKeyEngine(
checkpoint_path=checkpoint_path,
device=str(device),
img_size=2048,
use_refiner=use_refiner,
mixed_precision=mixed_precision,
model_precision=model_precision,
)
return (engine,)
requirements.txt
+2
View File
@@ -0,0 +1,2 @@
corridorkey @ git+https://github.com/nikopueringer/CorridorKey.git --no-deps
huggingface-hub