From 25ad75758ecf22a38566e4a5f42e490292bb384f Mon Sep 17 00:00:00 2001 From: pixelworld AI Date: Wed, 18 Mar 2026 08:25:41 -0500 Subject: [PATCH] Initial commit: ComfyUI nodes for CorridorKey greenscreen keyer 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) --- .gitignore | 3 + README.md | 50 +++++++++++++ __init__.py | 10 +++ nodes/__init__.py | 18 +++++ nodes/chroma_prepass.py | 116 ++++++++++++++++++++++++++++++ nodes/composite.py | 72 +++++++++++++++++++ nodes/greenscreen.py | 155 ++++++++++++++++++++++++++++++++++++++++ nodes/load_model.py | 117 ++++++++++++++++++++++++++++++ requirements.txt | 2 + 9 files changed, 543 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 __init__.py create mode 100644 nodes/__init__.py create mode 100644 nodes/chroma_prepass.py create mode 100644 nodes/composite.py create mode 100644 nodes/greenscreen.py create mode 100644 nodes/load_model.py create mode 100644 requirements.txt diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..22275c0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +__pycache__/ +*.pyc +CorridorKey/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..8fd7727 --- /dev/null +++ b/README.md @@ -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. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..099f8eb --- /dev/null +++ b/__init__.py @@ -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"] diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..4ff85ab --- /dev/null +++ b/nodes/__init__.py @@ -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", +} diff --git a/nodes/chroma_prepass.py b/nodes/chroma_prepass.py new file mode 100644 index 0000000..3a3eece --- /dev/null +++ b/nodes/chroma_prepass.py @@ -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,) diff --git a/nodes/composite.py b/nodes/composite.py new file mode 100644 index 0000000..bdf34ef --- /dev/null +++ b/nodes/composite.py @@ -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(),) diff --git a/nodes/greenscreen.py b/nodes/greenscreen.py new file mode 100644 index 0000000..4392b43 --- /dev/null +++ b/nodes/greenscreen.py @@ -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) diff --git a/nodes/load_model.py b/nodes/load_model.py new file mode 100644 index 0000000..cbc635b --- /dev/null +++ b/nodes/load_model.py @@ -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. /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,) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..2697d80 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +corridorkey @ git+https://github.com/nikopueringer/CorridorKey.git --no-deps +huggingface-hub