diff --git a/.gitignore b/.gitignore index f02b17f..355da57 100644 --- a/.gitignore +++ b/.gitignore @@ -3,7 +3,15 @@ __pycache__/ .pytest_cache/ .mypy_cache/ .ruff_cache/ +.coverage +.hypothesis/ .venv/ +.vscode/ +.idea/ +.DS_Store +Thumbs.db +models/*.pth dist/ build/ +pip-wheel-metadata/ *.egg-info/ diff --git a/README.md b/README.md index e5e0388..e61bed8 100644 --- a/README.md +++ b/README.md @@ -2,63 +2,72 @@ ## 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. +`ComfyUI-CorridorKey` is a ComfyUI custom node package that exposes a native `CorridorKey` inference node for four-pass CorridorKey processing 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: +`ComfyUI-CorridorKey` adapts the upstream `CorridorKey` inference flow for ComfyUI users who already build coarse alpha hints inside ComfyUI. Instead of running the upstream command-line wizard and instead of bundling extra mask-generation tools such as `VideoMaMa` or `GVM`, the main 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: +It then runs an upstream-style 4-channel model pass (RGB plus coarse alpha hint), resizes through the expected `2048x2048` inference resolution, 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 +- `fg`: straight foreground color +- `matte`: processed linear alpha +- `processed`: linear premultiplied RGB +- `QC`: sRGB checkerboard composite preview 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. +- Keep the real CorridorKey inference path available inside a normal ComfyUI graph. +- Preserve the upstream color math rules for sRGB and linear conversions. - 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 ship the CorridorKey model checkpoint. - 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 download models automatically. +- This project does not require network access for normal inference once dependencies and the checkpoint are installed. - 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` +- One ComfyUI node: + - `CorridorKey` +- Uses vendored upstream-style runtime pieces: + - `GreenFormer` + - `CNNRefinerModule` + - upstream color transfer math - 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: +- Requires a matching coarse alpha hint for each frame in a batch instead of silently reusing one mask across multiple frames +- Exposes upstream settings 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 +- Returns the four practical passes directly: + - `fg` + - `matte` + - `processed` + - `QC` +- Includes a bundled example workflow: + - `example_workflows/corridorkey_example_workflow.json` + - `example_workflows/PF0044-01_Clip-1_REC709_2K.mp4` + - `example_workflows/PF0044-01_Clip-1_REC709_2K_mask.mp4` - 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 +- Planned: a loader node that can expose a reusable cached engine object explicitly +- Planned: optional raw-alpha and raw-fg outputs in addition to the practical four-pass outputs +- Planned: optional direct support for a dedicated ComfyUI models search path ## 4) Requirements (Python version, OS, dependencies) @@ -66,7 +75,10 @@ Planned features: - OS: Windows, Linux, or macOS where ComfyUI runs - Runtime dependencies: - `torch` (normally already present in the ComfyUI Python environment) + - `torchvision` + - `timm` - `numpy` + - `opencv-python` - `Pillow` - Development dependencies: - `pytest` @@ -82,24 +94,45 @@ 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. +3. Download the CorridorKey model checkpoint and place it in `ComfyUI-CorridorKey/models/`. +4. Restart ComfyUI. Typical install flow: ```powershell cd ComfyUI\custom_nodes -git clone ComfyUI-CorridorKey +git clone https://github.com/SeanBRVFX/ComfyUI-CorridorKey.git cd ComfyUI-CorridorKey python -m pip install -r requirements.txt ``` +Before restarting ComfyUI, download the model checkpoint from: + +```text +https://huggingface.co/nikopueringer/CorridorKey_v1.0/resolve/main/CorridorKey_v1.0.pth +``` + +Then place that `.pth` file in: + +```text +ComfyUI-CorridorKey/models/ +``` + +Recommended filename: + +```text +ComfyUI-CorridorKey/models/CorridorKey.pth +``` + +The loader will also accept another `.pth` filename in `models/`, so the original `CorridorKey_v1.0.pth` filename works too if it is the only checkpoint in that folder. + 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 ComfyUI-CorridorKey +git clone https://github.com/SeanBRVFX/ComfyUI-CorridorKey.git ..\..\python_embeded\python.exe -m pip install -r .\ComfyUI-CorridorKey\requirements.txt ``` @@ -107,7 +140,8 @@ 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. +3. Add the checkpoint file under `models/`. +4. Restart ComfyUI. Development install: @@ -119,44 +153,46 @@ 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. +Repository publishing notes: + +- Repository URL: `https://github.com/SeanBRVFX/ComfyUI-CorridorKey` + +- This repo is already in the standard shape for both manual `git clone` installs and ComfyUI-Manager `git-clone` installs: + - the node files live at the repository root + - `requirements.txt` is in the repository root + - `README.md` is in the repository root + - `__init__.py` is in the repository root +- There is no special extra metadata file required inside this repository for basic ComfyUI-Manager compatibility. +- For ComfyUI-Manager discovery, the repository URL typically needs to be added to an external node list maintained by ComfyUI-Manager or another registry source. + ## 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` +- `gamma_space`: `sRGB` or `Linear` +- `despill_strength`: normalized float from `0.0` to `1.0` +- `refiner_strength`: float multiplier for the upstream refiner deltas +- `auto_despeckle`: enables the upstream-style connected-component cleanup pass +- `despeckle_size`: minimum island size in pixels for auto-despeckle 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: +Implementation notes: -- 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. +- This node uses the real upstream-style model path and color math, not the earlier heuristic approximation. - 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. +- The `mask` input is the upstream-style coarse alpha hint: rough, blurry, and usually slightly eroded works better than an expanded mask. +- For batched images, the mask batch must match the image batch. The node does not auto-repeat one mask across multiple frames. Environment variables for upstream tracking: @@ -172,43 +208,69 @@ Environment variables for upstream tracking: The currently reviewed upstream head SHA and its last observed check state are pinned in `corridor_key/upstream_sync.py`. +Environment variables for performance and CUDA behavior: + +- `CORRIDORKEY_PREFER_CHANNELS_LAST` + - Default: `1` + - Enables channels-last memory format for better CUDA throughput when supported +- `CORRIDORKEY_ENABLE_TF32` + - Default: `1` + - Allows TF32 on supported NVIDIA GPUs for better throughput with small quality tradeoffs typical for inference + ## 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. +2. Create a coarse alpha hint using existing ComfyUI nodes such as RMBG, SAM, Segment Anything, Florence-based tools, or manual masking nodes. +3. Keep that alpha hint rough rather than over-expanded. Slight erosion, softness, and blur are usually closer to what the upstream model was trained on. +4. Feed the `IMAGE` and coarse `MASK` into `CorridorKey`. +5. For batched images, provide one mask per frame. Do not rely on one static mask for a moving sequence. +6. Send the four pass outputs into downstream preview or save nodes. +7. Send the four pass outputs into your preferred existing saver or export nodes. +8. While the node runs, it emits brief progress text in ComfyUI and prints matching status lines in the ComfyUI console so long operations are visible. 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` +- `fg`: `IMAGE` +- `matte`: `MASK` - `processed`: `IMAGE` -- `comp`: `IMAGE` +- `QC`: `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. +- `Load Image` -> `RMBG` (or another segmentation node) -> `CorridorKey` -> your existing save/export nodes +- Branch the outputs into standard preview nodes and whichever saver nodes you already use for EXR, PNG, or compositing workflows. +- A bundled video example is included at `example_workflows/corridorkey_example_workflow.json`. It uses core `LoadVideo`, `GetVideoComponents`, `ImageToMask`, `CreateVideo`, and `SaveVideo` nodes plus `CorridorKey`. +- That bundled workflow is only a simple video-based test to confirm the node runs end-to-end inside ComfyUI. +- For best quality and behavior closer to the original CorridorKey workflow, prefer EXR image-sequence input and EXR image-sequence output instead of compressed video files. +- For EXR-focused ComfyUI workflows, `ComfyUI-CoCoTools_IO` is a recommended companion node pack: + - `https://github.com/Conor-Collins/ComfyUI-CoCoTools_IO` +- The bundled example references these included sample files by filename: + - `PF0028-05_Clip-5_REC709_2K_input.mp4` + - `PF0028-05_Clip-5_REC709_2K_mask.mp4` +- The bundled sample footage is sourced from ActionVFX practice footage: + - `https://www.actionvfx.com/practice-footage/mechanic-in-electrical-explosion/15806` + - Scene: `PF0028` + - Title: `Mechanic in Electrical Explosion` + - ActionVFX describes this footage as covered by its license agreement for commercial use, learning, teaching, demo reels, and similar use cases. Review the ActionVFX license terms directly if you plan to redistribute or publish derivative examples. +- The bundled mask example is intentionally simple: it was made in DaVinci Resolve as a basic rough mask with a very slight outward expand/erode adjustment plus blur, to keep it in coarse alpha-hint territory rather than final roto. +- An observed end-to-end runtime for the bundled video test was `484` seconds on this reported machine: + - `128 GB` RAM + - `RTX 4070 Ti Super` + - `16 GB` VRAM + - This is only a real-world reference point for the bundled sample, not a formal benchmark. Actual runtime will vary with driver state, other loaded ComfyUI models, codec behavior, and source resolution or duration. +- If your local `LoadVideo` node does not automatically resolve files from `example_workflows`, use the same filenames and reselect the bundled files once in the node UI. There is no standalone CLI in this custom node package by design. @@ -231,62 +293,65 @@ IMAGE + coarse MASK input validation | v - mask normalization + batch normalization | v - optional gamma handling + resize to 2048x2048 | v - image luminance + edge map + 4-channel model inference | v - refiner + despeckle + upstream refiner delta scaling | v - refined alpha MASK + linear alpha + straight FG | v - despill + output assembly - | | | - v v v - FG Processed Comp + auto-despeckle + despill + | + v + output assembly + | | | | + v v v v + FG Matte Processed QC ``` ### 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 + - ComfyUI-facing `CorridorKey` node - `corridor_key/config.py` - Dataclass for validated runtime settings +- `corridor_key/color_utils.py` + - Vendored upstream sRGB/linear conversions, despill, composite, and matte cleanup helpers +- `corridor_key/model_transformer.py` + - Vendored upstream `GreenFormer` and `CNNRefinerModule` +- `corridor_key/engine.py` + - Vendored and adapted upstream inference engine with 2048 resize handling - `corridor_key/tensor_ops.py` - - Shared tensor validation, mask normalization, color helpers, and small image-processing helpers + - Tensor validation and tensor/NumPy batch conversion helpers - `corridor_key/processor.py` - - Pure processing pipeline for edge-aware matte refinement, despill, despeckle, and output assembly + - Batch bridge between ComfyUI tensors and the upstream-style engine - `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 + - Unit tests for configuration, color math, and sync policy helpers ### 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. +3. The processor validates the image and mask tensors and aligns the batch. +4. The engine resizes each frame to the fixed `2048x2048` inference size, converts gamma as needed, and runs the 4-channel model. +5. The engine applies upstream-style auto-despeckle, despill, linear premultiplication, and checkerboard compositing. +6. The processor returns: + - `fg` + - `matte` + - `processed` + - `QC` +7. The four outputs are then available to downstream preview, save, or export nodes already present in your ComfyUI setup. +8. In parallel, an optional background task may query GitHub and log if a newer verified upstream commit is available. ## 9) Repository structure (tree) @@ -298,9 +363,18 @@ ComfyUI-CorridorKey/ |-- pyproject.toml |-- __init__.py |-- nodes.py +|-- example_workflows/ +| `-- corridorkey_example_workflow.json +| |-- PF0044-01_Clip-1_REC709_2K.mp4 +| `-- PF0044-01_Clip-1_REC709_2K_mask.mp4 +|-- models/ +| `-- CorridorKey.pth (user-supplied, not committed) |-- corridor_key/ | |-- __init__.py | |-- config.py +| |-- color_utils.py +| |-- engine.py +| |-- model_transformer.py | |-- tensor_ops.py | |-- processor.py | `-- upstream_sync.py @@ -329,15 +403,14 @@ When changing behavior: ### testing strategy (pytest) -- Test pure processing code in `corridor_key/processor.py` -- Avoid ComfyUI runtime dependencies in the test suite where possible +- Test configuration and utility code without requiring the checkpoint where possible +- Avoid relying on a full model load in unit tests unless an explicit integration test is added later - Cover: - - mask broadcasting - parameter validation - - output shapes - - value range clamping - - output mode selection - - despeckle and despill edge cases + - color transfer function behavior + - batch shape handling + - four-pass output shapes + - despill and auto-despeckle edge cases - verified-upstream selection policy for mocked GitHub metadata ### lint/format rules (ruff + black or ruff format) @@ -357,10 +430,12 @@ When changing behavior: - 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 +- Runtime status output is intentionally concise: + - model-loading errors are raised with explicit setup guidance + - the node emits short progress text while loading the model and processing frames + - the same progress is printed to the ComfyUI console for visibility during long runs + - a module-level logger may still 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) @@ -370,7 +445,7 @@ When changing behavior: - 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 +- Inference reads only the user-supplied checkpoint under the project directory - 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 @@ -380,7 +455,8 @@ 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. +3. Add the checkpoint file under `models/`. +4. Restart ComfyUI. There is no separate server, container, or packaged service in this repository. @@ -388,39 +464,52 @@ There is no separate server, container, or packaged service in this repository. Q: Why does this node require a mask input? -A: This package intentionally reuses existing ComfyUI segmentation and masking nodes instead of duplicating them. +A: The input is the upstream-style coarse alpha hint. 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? +Q: Can I use one mask for a whole image batch or video batch? -A: This implementation provides a native, deterministic refinement pass inspired by the same workflow goal, without bundling upstream research weights or optional external pipelines. +A: No. For batched inputs, provide a matching mask for each frame. CorridorKey expects the coarse alpha hint to describe the current frame, even if that hint is rough and slightly eroded. -Q: My mask looks too soft. +Q: The node says the checkpoint is missing. -A: Lower `feather_radius`, lower `preserve_core`, raise `edge_focus`, or raise `refiner_strength`. +A: Place the `CorridorKey.pth` model file in `ComfyUI-CorridorKey/models/` and restart ComfyUI. -Q: My mask looks too harsh. +Checkpoint download URL: -A: Raise `feather_radius` slightly or lower `threshold`. +`https://huggingface.co/nikopueringer/CorridorKey_v1.0/resolve/main/CorridorKey_v1.0.pth` 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. +A: `Processed` is returned as premultiplied RGB. Use your existing ComfyUI saver/export nodes to pack or write it the way you want, including workflows that combine `processed` with `matte` downstream. + +Q: Why does the bundled example workflow use video if EXR is better? + +A: The bundled workflow is only a simple native-node test. For production work and behavior closer to the original CorridorKey workflow, use EXR image sequences for input and output. If you want EXR-focused ComfyUI IO, `ComfyUI-CoCoTools_IO` is a strong companion option: `https://github.com/Conor-Collins/ComfyUI-CoCoTools_IO` 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: Why does strong despill sometimes look purple? + +A: Higher despill values remove excess green and rebalance color into red and blue, which can push heavy spill areas toward magenta or purple. If that happens, lower `despill_strength` and do the remaining color cleanup later in comp if needed. + 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. +Q: Can this work on 16GB consumer GPUs? + +A: That is a design goal, but the inference path stays fixed at `2048x2048` because that is what the model is trained for. Practical support on 16GB-class GPUs should come from memory and throughput optimization at the fixed size, not from silently changing the inference resolution. + ## 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 +- Prioritize runtime performance first: reduce per-frame overhead, minimize transfers, and keep throughput strong on production-sized sequences +- Improve practical support for consumer GPUs, especially 16GB-class cards, while keeping the fixed `2048x2048` inference path intact +- Add a loader node that can expose a reusable cached engine object explicitly +- Add optional raw-alpha and raw-fg outputs in addition to the practical four-pass outputs +- Add optional direct support for a dedicated ComfyUI models search path +- Expand the bundled example workflows and keep their real-world benchmark notes current as the node evolves ## 16) Contributing (branching, commit style, PR checklist) @@ -432,7 +521,7 @@ Branching: Commit style: - Prefer imperative commit messages -- Example: `Add edge-aware CorridorKey refinement node` +- Example: `Port upstream CorridorKey runtime into ComfyUI` PR checklist: @@ -443,6 +532,19 @@ PR checklist: - Run `mypy corridor_key` - Confirm the node loads in ComfyUI - Keep the public node contract backward compatible unless documented +- Performance and cleanup contributions are explicitly welcome. This wrapper was prototyped quickly, and contributors with deeper ComfyUI, PyTorch, or optimization experience can improve it further. + +Release checklist: + +- Confirm `.gitignore` excludes caches and local editor files before the first commit +- Ensure `requirements.txt` matches runtime imports +- Ensure `pyproject.toml` version is correct for the release +- Ensure `README.md` installation instructions match the actual repo URL and current node name +- Make sure `__pycache__` folders are not staged +- Create a clean commit on `main` +- Tag the release, for example `v0.1.0` +- Push both `main` and the tag to GitHub +- Verify a fresh `git clone` into `ComfyUI/custom_nodes` works with `python -m pip install -r requirements.txt` ## 17) License placeholder diff --git a/__init__.py b/__init__.py index c5e6f1e..f1e08f7 100644 --- a/__init__.py +++ b/__init__.py @@ -2,7 +2,7 @@ @author: local @title: ComfyUI-CorridorKey @nickname: CorridorKey -@description: ComfyUI node for CorridorKey-style coarse mask refinement. +@description: ComfyUI node for CorridorKey inference. """ from __future__ import annotations diff --git a/corridor_key/color_utils.py b/corridor_key/color_utils.py new file mode 100644 index 0000000..f183cfb --- /dev/null +++ b/corridor_key/color_utils.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import numpy as np +import torch + + +def _is_tensor(value: np.ndarray | torch.Tensor) -> bool: + return isinstance(value, torch.Tensor) + + +def linear_to_srgb(value: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: + if _is_tensor(value): + value = value.clamp(min=0.0) + mask = value <= 0.0031308 + return torch.where(mask, value * 12.92, 1.055 * torch.pow(value, 1.0 / 2.4) - 0.055) + + value = np.clip(value, 0.0, None) + mask = value <= 0.0031308 + return np.where(mask, value * 12.92, 1.055 * np.power(value, 1.0 / 2.4) - 0.055) + + +def srgb_to_linear(value: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: + if _is_tensor(value): + value = value.clamp(min=0.0) + mask = value <= 0.04045 + return torch.where(mask, value / 12.92, torch.pow((value + 0.055) / 1.055, 2.4)) + + value = np.clip(value, 0.0, None) + mask = value <= 0.04045 + return np.where(mask, value / 12.92, np.power((value + 0.055) / 1.055, 2.4)) + + +def premultiply( + foreground: np.ndarray | torch.Tensor, + alpha: np.ndarray | torch.Tensor, +) -> np.ndarray | torch.Tensor: + return foreground * alpha + + +def composite_straight( + foreground: np.ndarray | torch.Tensor, + background: np.ndarray | torch.Tensor, + alpha: np.ndarray | torch.Tensor, +) -> np.ndarray | torch.Tensor: + return foreground * alpha + background * (1.0 - alpha) + + +def composite_premul( + foreground: np.ndarray | torch.Tensor, + background: np.ndarray | torch.Tensor, + alpha: np.ndarray | torch.Tensor, +) -> np.ndarray | torch.Tensor: + return foreground + background * (1.0 - alpha) + + +def despill( + image: np.ndarray | torch.Tensor, + green_limit_mode: str = "average", + strength: float = 1.0, +) -> np.ndarray | torch.Tensor: + if strength <= 0.0: + return image + + if _is_tensor(image): + red = image[..., 0] + green = image[..., 1] + blue = image[..., 2] + + limit = torch.max(red, blue) if green_limit_mode == "max" else (red + blue) / 2.0 + spill_amount = torch.clamp(green - limit, min=0.0) + + despilled = torch.stack( + [ + red + (spill_amount * 0.5), + green - spill_amount, + blue + (spill_amount * 0.5), + ], + dim=-1, + ) + if strength < 1.0: + return image * (1.0 - strength) + despilled * strength + return despilled + + red = image[..., 0] + green = image[..., 1] + blue = image[..., 2] + + limit = np.maximum(red, blue) if green_limit_mode == "max" else (red + blue) / 2.0 + spill_amount = np.maximum(green - limit, 0.0) + + despilled_np = np.stack( + [ + red + (spill_amount * 0.5), + green - spill_amount, + blue + (spill_amount * 0.5), + ], + axis=-1, + ) + if strength < 1.0: + return image * (1.0 - strength) + despilled_np * strength + return despilled_np + + +def clean_matte( + alpha_np: np.ndarray, + area_threshold: int = 300, + dilation: int = 15, + blur_size: int = 5, +) -> np.ndarray: + import cv2 + + is_3d = alpha_np.ndim == 3 + working = alpha_np[:, :, 0] if is_3d else alpha_np + + mask_8u = (working > 0.5).astype(np.uint8) * 255 + num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask_8u, connectivity=8) + + cleaned_mask = np.zeros_like(mask_8u) + for label_index in range(1, num_labels): + if stats[label_index, cv2.CC_STAT_AREA] >= area_threshold: + cleaned_mask[labels == label_index] = 255 + + if dilation > 0: + kernel_size = int(dilation * 2 + 1) + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + cleaned_mask = cv2.dilate(cleaned_mask, kernel) + + if blur_size > 0: + blur_kernel = int(blur_size * 2 + 1) + cleaned_mask = cv2.GaussianBlur(cleaned_mask, (blur_kernel, blur_kernel), 0) + + safe_zone = cleaned_mask.astype(np.float32) / 255.0 + result = working * safe_zone + if is_3d: + return result[:, :, np.newaxis] + return result + + +def create_checkerboard( + width: int, + height: int, + checker_size: int = 64, + color1: float = 0.2, + color2: float = 0.4, +) -> np.ndarray: + x_tiles = np.arange(width) // checker_size + y_tiles = np.arange(height) // checker_size + x_grid, y_grid = np.meshgrid(x_tiles, y_tiles) + checker = (x_grid + y_grid) % 2 + background = np.where(checker == 0, color1, color2).astype(np.float32) + return np.stack([background, background, background], axis=-1) diff --git a/corridor_key/config.py b/corridor_key/config.py index 372f93e..d11f0ab 100644 --- a/corridor_key/config.py +++ b/corridor_key/config.py @@ -6,45 +6,26 @@ 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 not 0.0 <= self.refiner_strength <= 4.0: + raise ValueError("refiner_strength must be between 0.0 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.") + if not 0 <= self.despeckle_size <= 4096: + raise ValueError("despeckle_size must be between 0 and 4096.") @property - def transition_gain(self) -> float: - return 8.0 + (self.edge_focus * 16.0) - - @property - def despill_factor(self) -> float: - return self.despill_strength + def input_is_linear(self) -> bool: + return self.gamma_space == "Linear" @property def despeckle_enabled(self) -> bool: diff --git a/corridor_key/engine.py b/corridor_key/engine.py new file mode 100644 index 0000000..9ebc3e6 --- /dev/null +++ b/corridor_key/engine.py @@ -0,0 +1,239 @@ +from __future__ import annotations + +import os +from contextlib import nullcontext +from pathlib import Path + +import numpy as np +import torch + +from . import color_utils as cu +from .model_transformer import GreenFormer + +PROJECT_ROOT = Path(__file__).resolve().parent.parent +MODELS_DIR = PROJECT_ROOT / "models" +DEFAULT_CHECKPOINT_NAME = "CorridorKey.pth" +CHECKPOINT_DOWNLOAD_URL = "https://huggingface.co/nikopueringer/CorridorKey_v1.0/resolve/main/CorridorKey_v1.0.pth" +_ENGINE_CACHE: dict[tuple[str, str, int, bool], "CorridorKeyEngine"] = {} + + +def _import_cv2(): + os.environ.setdefault("OPENCV_IO_ENABLE_OPENEXR", "1") + try: + import cv2 + except ImportError as exc: + raise ImportError( + "opencv-python is required for CorridorKey. Install requirements.txt in the ComfyUI Python environment." + ) from exc + return cv2 + + +def _parse_bool_env(name: str, default: bool) -> bool: + raw = os.getenv(name) + if raw is None: + return default + return raw.strip().lower() not in {"0", "false", "no", "off"} + + +def _prefer_channels_last() -> bool: + return _parse_bool_env("CORRIDORKEY_PREFER_CHANNELS_LAST", True) + + +def _enable_tf32() -> bool: + return _parse_bool_env("CORRIDORKEY_ENABLE_TF32", True) + + +def _configure_torch_for_inference(device_type: str) -> None: + if device_type != "cuda": + return + if hasattr(torch.backends, "cudnn"): + torch.backends.cudnn.benchmark = True + if hasattr(torch.backends.cudnn, "allow_tf32"): + torch.backends.cudnn.allow_tf32 = _enable_tf32() + if hasattr(torch.backends, "cuda") and hasattr(torch.backends.cuda, "matmul"): + torch.backends.cuda.matmul.allow_tf32 = _enable_tf32() + + +def resolve_checkpoint_path() -> Path: + explicit = MODELS_DIR / DEFAULT_CHECKPOINT_NAME + if explicit.is_file(): + return explicit + + if not MODELS_DIR.exists(): + raise FileNotFoundError( + f"CorridorKey model directory is missing: {MODELS_DIR}. " + "Create it and place CorridorKey.pth inside. " + f"Download URL: {CHECKPOINT_DOWNLOAD_URL}" + ) + + candidates = sorted(MODELS_DIR.glob("*.pth")) + if not candidates: + raise FileNotFoundError( + f"No CorridorKey model file found in {MODELS_DIR}. " + "Download the model checkpoint and place it there. " + f"Download URL: {CHECKPOINT_DOWNLOAD_URL}" + ) + return candidates[0] + + +class CorridorKeyEngine: + def __init__( + self, + checkpoint_path: Path, + device: str | None = None, + img_size: int = 2048, + use_refiner: bool = True, + ) -> None: + self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) + self.img_size = img_size + self.checkpoint_path = Path(checkpoint_path) + self.use_refiner = use_refiner + self.channels_last = self.device.type == "cuda" and _prefer_channels_last() + + self.mean = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(1, 1, 3) + self.std = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(1, 1, 3) + _configure_torch_for_inference(self.device.type) + self.model = self._load_model() + + def _load_model(self) -> GreenFormer: + model = GreenFormer( + encoder_name="hiera_base_plus_224.mae_in1k_ft_in1k", + img_size=self.img_size, + use_refiner=self.use_refiner, + ) + model = model.to(self.device) + if self.channels_last: + model = model.to(memory_format=torch.channels_last) + model.eval() + + if not self.checkpoint_path.is_file(): + raise FileNotFoundError( + f"CorridorKey model file not found: {self.checkpoint_path}. " + "Place CorridorKey.pth in the local models directory. " + f"Download URL: {CHECKPOINT_DOWNLOAD_URL}" + ) + + checkpoint = torch.load(str(self.checkpoint_path), map_location=self.device) + state_dict = checkpoint.get("state_dict", checkpoint) + model.load_checkpoint(state_dict) + return model + + def _autocast_context(self): + if self.device.type == "cuda": + return torch.autocast(device_type="cuda", dtype=torch.float16) + return nullcontext() + + @torch.no_grad() + def process_frame( + self, + image: np.ndarray, + mask_linear: np.ndarray, + refiner_scale: float = 1.0, + input_is_linear: bool = False, + fg_is_straight: bool = True, + despill_strength: float = 1.0, + auto_despeckle: bool = True, + despeckle_size: int = 400, + ) -> dict[str, np.ndarray]: + cv2 = _import_cv2() + + if image.dtype == np.uint8: + image = image.astype(np.float32) / 255.0 + else: + image = image.astype(np.float32) + + if mask_linear.dtype == np.uint8: + mask_linear = mask_linear.astype(np.float32) / 255.0 + elif mask_linear.dtype == np.uint16: + mask_linear = mask_linear.astype(np.float32) / 65535.0 + else: + mask_linear = mask_linear.astype(np.float32) + + height, width = image.shape[:2] + + if mask_linear.ndim == 2: + mask_linear = mask_linear[:, :, np.newaxis] + + if input_is_linear: + img_resized_linear = cv2.resize(image, (self.img_size, self.img_size), interpolation=cv2.INTER_LINEAR) + img_resized = cu.linear_to_srgb(img_resized_linear) + else: + img_resized = cv2.resize(image, (self.img_size, self.img_size), interpolation=cv2.INTER_LINEAR) + + mask_resized = cv2.resize(mask_linear, (self.img_size, self.img_size), interpolation=cv2.INTER_LINEAR) + if mask_resized.ndim == 2: + mask_resized = mask_resized[:, :, np.newaxis] + + img_norm = (img_resized - self.mean) / self.std + input_np = np.concatenate([img_norm, mask_resized], axis=-1) + input_tensor = torch.from_numpy(input_np.transpose((2, 0, 1))).float().unsqueeze(0).to(self.device) + if self.channels_last: + input_tensor = input_tensor.to(memory_format=torch.channels_last) + + hook_handle = None + if refiner_scale != 1.0 and self.model.refiner is not None: + def scale_hook(_module, _inputs, output): + return output * refiner_scale + + hook_handle = self.model.refiner.register_forward_hook(scale_hook) + + with self._autocast_context(): + output = self.model(input_tensor) + + if hook_handle is not None: + hook_handle.remove() + + pred_alpha = output["alpha"] + pred_fg = output["fg"] + + resized_alpha = pred_alpha[0].permute(1, 2, 0).cpu().numpy() + resized_fg = pred_fg[0].permute(1, 2, 0).cpu().numpy() + resized_alpha = cv2.resize(resized_alpha, (width, height), interpolation=cv2.INTER_LANCZOS4) + resized_fg = cv2.resize(resized_fg, (width, height), interpolation=cv2.INTER_LANCZOS4) + + if resized_alpha.ndim == 2: + resized_alpha = resized_alpha[:, :, np.newaxis] + + if auto_despeckle: + processed_alpha = cu.clean_matte(resized_alpha, area_threshold=despeckle_size, dilation=25, blur_size=5) + else: + processed_alpha = resized_alpha + + fg_despilled = cu.despill(resized_fg, green_limit_mode="average", strength=despill_strength) + fg_despilled_linear = cu.srgb_to_linear(fg_despilled) + fg_premul_linear = cu.premultiply(fg_despilled_linear, processed_alpha) + processed_rgba = np.concatenate([fg_premul_linear, processed_alpha], axis=-1) + + checkerboard_srgb = cu.create_checkerboard(width, height, checker_size=128, color1=0.15, color2=0.55) + checkerboard_linear = cu.srgb_to_linear(checkerboard_srgb) + if fg_is_straight: + comp_linear = cu.composite_straight(fg_despilled_linear, checkerboard_linear, processed_alpha) + else: + comp_linear = cu.composite_premul(fg_despilled_linear, checkerboard_linear, processed_alpha) + comp_srgb = cu.linear_to_srgb(comp_linear) + + return { + "fg": np.clip(resized_fg.astype(np.float32), 0.0, 1.0), + "raw_alpha": np.clip(resized_alpha.astype(np.float32), 0.0, 1.0), + "matte": np.clip(processed_alpha.astype(np.float32), 0.0, 1.0), + "processed": np.clip(fg_premul_linear.astype(np.float32), 0.0, 1.0), + "processed_rgba": np.clip(processed_rgba.astype(np.float32), 0.0, 1.0), + "comp": np.clip(comp_srgb.astype(np.float32), 0.0, 1.0), + } + + +def get_cached_engine(device: str | None = None, img_size: int = 2048, use_refiner: bool = True) -> CorridorKeyEngine: + checkpoint_path = resolve_checkpoint_path() + cache_key = (str(checkpoint_path), device or "", img_size, use_refiner) + cached = _ENGINE_CACHE.get(cache_key) + if cached is not None: + return cached + + engine = CorridorKeyEngine( + checkpoint_path=checkpoint_path, + device=device, + img_size=img_size, + use_refiner=use_refiner, + ) + _ENGINE_CACHE[cache_key] = engine + return engine diff --git a/corridor_key/model_transformer.py b/corridor_key/model_transformer.py new file mode 100644 index 0000000..48066a7 --- /dev/null +++ b/corridor_key/model_transformer.py @@ -0,0 +1,234 @@ +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class MLP(nn.Module): + def __init__(self, input_dim: int = 2048, embed_dim: int = 768) -> None: + super().__init__() + self.proj = nn.Linear(input_dim, embed_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.proj(x) + + +class DecoderHead(nn.Module): + def __init__( + self, + feature_channels: list[int] | None = None, + embedding_dim: int = 256, + output_dim: int = 1, + ) -> None: + super().__init__() + if feature_channels is None: + feature_channels = [112, 224, 448, 896] + + self.linear_c4 = MLP(input_dim=feature_channels[3], embed_dim=embedding_dim) + self.linear_c3 = MLP(input_dim=feature_channels[2], embed_dim=embedding_dim) + self.linear_c2 = MLP(input_dim=feature_channels[1], embed_dim=embedding_dim) + self.linear_c1 = MLP(input_dim=feature_channels[0], embed_dim=embedding_dim) + + self.linear_fuse = nn.Conv2d(embedding_dim * 4, embedding_dim, kernel_size=1, bias=False) + self.bn = nn.BatchNorm2d(embedding_dim) + self.relu = nn.ReLU(inplace=True) + self.dropout = nn.Dropout(0.1) + self.classifier = nn.Conv2d(embedding_dim, output_dim, kernel_size=1) + + def forward(self, features: list[torch.Tensor]) -> torch.Tensor: + c1, c2, c3, c4 = features + batch_size = c4.shape[0] + + c4_out = self.linear_c4(c4.flatten(2).transpose(1, 2)).transpose(1, 2).view( + batch_size, + -1, + c4.shape[2], + c4.shape[3], + ) + c4_out = F.interpolate(c4_out, size=c1.shape[2:], mode="bilinear", align_corners=False) + + c3_out = self.linear_c3(c3.flatten(2).transpose(1, 2)).transpose(1, 2).view( + batch_size, + -1, + c3.shape[2], + c3.shape[3], + ) + c3_out = F.interpolate(c3_out, size=c1.shape[2:], mode="bilinear", align_corners=False) + + c2_out = self.linear_c2(c2.flatten(2).transpose(1, 2)).transpose(1, 2).view( + batch_size, + -1, + c2.shape[2], + c2.shape[3], + ) + c2_out = F.interpolate(c2_out, size=c1.shape[2:], mode="bilinear", align_corners=False) + + c1_out = self.linear_c1(c1.flatten(2).transpose(1, 2)).transpose(1, 2).view( + batch_size, + -1, + c1.shape[2], + c1.shape[3], + ) + + fused = self.linear_fuse(torch.cat([c4_out, c3_out, c2_out, c1_out], dim=1)) + fused = self.bn(fused) + fused = self.relu(fused) + fused = self.dropout(fused) + return self.classifier(fused) + + +class RefinerBlock(nn.Module): + def __init__(self, channels: int, dilation: int = 1) -> None: + super().__init__() + self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=dilation, dilation=dilation) + self.gn1 = nn.GroupNorm(8, channels) + self.relu = nn.ReLU(inplace=True) + self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=dilation, dilation=dilation) + self.gn2 = nn.GroupNorm(8, channels) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + out = self.conv1(x) + out = self.gn1(out) + out = self.relu(out) + out = self.conv2(out) + out = self.gn2(out) + out = out + residual + return self.relu(out) + + +class CNNRefinerModule(nn.Module): + def __init__(self, in_channels: int = 7, hidden_channels: int = 64, out_channels: int = 4) -> None: + super().__init__() + self.stem = nn.Sequential( + nn.Conv2d(in_channels, hidden_channels, kernel_size=3, padding=1), + nn.GroupNorm(8, hidden_channels), + nn.ReLU(inplace=True), + ) + self.res1 = RefinerBlock(hidden_channels, dilation=1) + self.res2 = RefinerBlock(hidden_channels, dilation=2) + self.res3 = RefinerBlock(hidden_channels, dilation=4) + self.res4 = RefinerBlock(hidden_channels, dilation=8) + self.final = nn.Conv2d(hidden_channels, out_channels, kernel_size=1) + + nn.init.normal_(self.final.weight, mean=0.0, std=1e-3) + nn.init.constant_(self.final.bias, 0) + + def forward(self, image: torch.Tensor, coarse_pred: torch.Tensor) -> torch.Tensor: + x = torch.cat([image, coarse_pred], dim=1) + x = self.stem(x) + x = self.res1(x) + x = self.res2(x) + x = self.res3(x) + x = self.res4(x) + return self.final(x) * 10.0 + + +class GreenFormer(nn.Module): + def __init__( + self, + encoder_name: str = "hiera_base_plus_224.mae_in1k_ft_in1k", + in_channels: int = 4, + img_size: int = 2048, + use_refiner: bool = True, + ) -> None: + super().__init__() + try: + import timm + except ImportError as exc: + raise ImportError( + "timm is required for CorridorKey. Install requirements.txt in the ComfyUI Python environment." + ) from exc + + self.encoder = timm.create_model(encoder_name, pretrained=False, features_only=True, img_size=img_size) + if in_channels != 3: + self._patch_input_layer(in_channels) + + try: + feature_channels = self.encoder.feature_info.channels() + except (AttributeError, TypeError): + feature_channels = [112, 224, 448, 896] + + self.alpha_decoder = DecoderHead(feature_channels, 256, output_dim=1) + self.fg_decoder = DecoderHead(feature_channels, 256, output_dim=3) + + self.use_refiner = use_refiner + self.refiner = CNNRefinerModule(in_channels=7, hidden_channels=64, out_channels=4) if use_refiner else None + + def _patch_input_layer(self, in_channels: int) -> None: + try: + patch_embed = self.encoder.model.patch_embed.proj + except AttributeError: + patch_embed = self.encoder.patch_embed.proj + + weight = patch_embed.weight.data + bias = patch_embed.bias.data if patch_embed.bias is not None else None + out_channels, _, kernel_h, kernel_w = weight.shape + + new_conv = nn.Conv2d( + in_channels, + out_channels, + kernel_size=(kernel_h, kernel_w), + stride=patch_embed.stride, + padding=patch_embed.padding, + bias=(bias is not None), + ) + new_conv.weight.data[:, :3, :, :] = weight + new_conv.weight.data[:, 3:, :, :] = 0.0 + if bias is not None: + new_conv.bias.data = bias + + try: + self.encoder.model.patch_embed.proj = new_conv + except AttributeError: + self.encoder.patch_embed.proj = new_conv + + def load_checkpoint(self, state_dict: dict[str, torch.Tensor]) -> tuple[list[str], list[str]]: + new_state_dict: dict[str, torch.Tensor] = {} + model_state = self.state_dict() + + for key, value in state_dict.items(): + if key.startswith("_orig_mod."): + key = key[10:] + + if "pos_embed" in key and key in model_state and value.shape != model_state[key].shape: + src_tokens = value.shape[1] + dst_tokens = model_state[key].shape[1] + channels = value.shape[2] + src_grid = int(math.sqrt(src_tokens)) + dst_grid = int(math.sqrt(dst_tokens)) + reshaped = value.permute(0, 2, 1).view(1, channels, src_grid, src_grid) + resized = F.interpolate(reshaped, size=(dst_grid, dst_grid), mode="bicubic", align_corners=False) + value = resized.flatten(2).transpose(1, 2) + + new_state_dict[key] = value + + missing, unexpected = self.load_state_dict(new_state_dict, strict=False) + return list(missing), list(unexpected) + + def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]: + input_size = x.shape[2:] + features = self.encoder(x) + + alpha_logits = self.alpha_decoder(features) + fg_logits = self.fg_decoder(features) + + alpha_logits_up = F.interpolate(alpha_logits, size=input_size, mode="bilinear", align_corners=False) + fg_logits_up = F.interpolate(fg_logits, size=input_size, mode="bilinear", align_corners=False) + + alpha_coarse = torch.sigmoid(alpha_logits_up) + fg_coarse = torch.sigmoid(fg_logits_up) + rgb = x[:, :3, :, :] + coarse_pred = torch.cat([alpha_coarse, fg_coarse], dim=1) + + if self.use_refiner and self.refiner is not None: + delta_logits = self.refiner(rgb, coarse_pred) + else: + delta_logits = torch.zeros_like(coarse_pred) + + alpha_final = torch.sigmoid(alpha_logits_up + delta_logits[:, 0:1]) + fg_final = torch.sigmoid(fg_logits_up + delta_logits[:, 1:4]) + return {"alpha": alpha_final, "fg": fg_final} diff --git a/corridor_key/processor.py b/corridor_key/processor.py index 7c88876..0afa004 100644 --- a/corridor_key/processor.py +++ b/corridor_key/processor.py @@ -1,133 +1,84 @@ from __future__ import annotations import logging +from typing import Callable import torch from .config import CorridorKeySettings +from .engine import get_cached_engine from .tensor_ops import ( - build_checkerboard, - clamp_unit_interval, - compute_edge_map, - convert_gamma_space, - despeckle_binary_mask, + batch_to_numpy, ensure_image_tensor, ensure_mask_batch, - gaussian_blur_map, - mask_to_image, + stack_mask_frames, + stack_rgb_frames, ) 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 __init__(self, device: str | None = None) -> None: + self._device = device def refine( self, image: torch.Tensor, mask: torch.Tensor, settings: CorridorKeySettings, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + progress_callback: Callable[[str, int, int], None] | None = None, + ) -> tuple[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], + batch_size=image_batch.shape[0], + height=image_batch.shape[1], + width=image_batch.shape[2], ) - # Mini title: compute the soft starting matte. - softened_mask = gaussian_blur_map(mask_batch, settings.feather_radius) + total_frames = int(image_batch.shape[0]) + if progress_callback is not None: + progress_callback("Loading CorridorKey model...", 0, total_frames) - # Mini title: derive edge guidance from image luminance. - edge_map = compute_edge_map(working_image) - edge_gain = 1.0 + (edge_map * settings.edge_focus) + engine = get_cached_engine(device=self._device) - centered_mask = ( - (softened_mask - settings.threshold) - * settings.transition_gain - * settings.refiner_strength - ) - alpha = torch.sigmoid(centered_mask * edge_gain) + fg_frames: list = [] + matte_frames: list = [] + processed_frames: list = [] + comp_frames: list = [] - 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) + image_frames = batch_to_numpy(image_batch) + mask_frames = batch_to_numpy(mask_batch) - # 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, + for frame_index, (image_frame, mask_frame) in enumerate(zip(image_frames, mask_frames, strict=True), start=1): + if progress_callback is not None: + progress_callback(f"Processing frame {frame_index}/{total_frames}...", frame_index - 1, total_frames) + + result = engine.process_frame( + image=image_frame, + mask_linear=mask_frame, + refiner_scale=settings.refiner_strength, + input_is_linear=settings.input_is_linear, + despill_strength=settings.despill_strength, + auto_despeckle=settings.despeckle_enabled, + despeckle_size=settings.despeckle_size, ) - alpha = torch.where(support_mask > 0.0, alpha, torch.zeros_like(alpha)) + fg_frames.append(result["fg"]) + matte_frames.append(result["matte"]) + processed_frames.append(result["processed"]) + comp_frames.append(result["comp"]) + LOGGER.debug("CorridorKey processed frame %s", frame_index) - alpha = clamp_unit_interval(alpha) + if progress_callback is not None: + progress_callback("CorridorKey complete.", total_frames, total_frames) - # 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, + return ( + stack_rgb_frames(fg_frames), + stack_mask_frames(matte_frames), + stack_rgb_frames(processed_frames), + stack_rgb_frames(comp_frames), ) - 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 diff --git a/corridor_key/tensor_ops.py b/corridor_key/tensor_ops.py index 1b99901..b79f896 100644 --- a/corridor_key/tensor_ops.py +++ b/corridor_key/tensor_ops.py @@ -1,177 +1,67 @@ from __future__ import annotations -from collections import deque +from typing import Iterable 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: + if min(image.shape[0], image.shape[1], 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: +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: + elif normalized.ndim != 3: 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.") + # Mini title: enforce per-frame alpha hints + if normalized.shape[0] != batch_size: + raise ValueError("mask batch dimension must match the image batch exactly.") 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 batch_to_numpy(batch: torch.Tensor) -> list[np.ndarray]: + return [frame.detach().cpu().numpy().astype(np.float32) for frame in batch] -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 stack_rgb_frames(frames: Iterable[np.ndarray]) -> torch.Tensor: + materialized = [np.clip(frame.astype(np.float32), 0.0, 1.0) for frame in frames] + if not materialized: + raise ValueError("No RGB frames were produced.") + return torch.from_numpy(np.stack(materialized, axis=0)) -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) +def stack_mask_frames(frames: Iterable[np.ndarray]) -> torch.Tensor: + materialized: list[np.ndarray] = [] + for frame in frames: + working = frame.astype(np.float32) + if working.ndim == 3 and working.shape[-1] == 1: + working = working[:, :, 0] + materialized.append(np.clip(working, 0.0, 1.0)) + if not materialized: + raise ValueError("No mask frames were produced.") + return torch.from_numpy(np.stack(materialized, axis=0)) diff --git a/example_workflows/PF0028-05_Clip-5_REC709_2K_input.mp4 b/example_workflows/PF0028-05_Clip-5_REC709_2K_input.mp4 new file mode 100644 index 0000000..b33a84d Binary files /dev/null and b/example_workflows/PF0028-05_Clip-5_REC709_2K_input.mp4 differ diff --git a/example_workflows/PF0028-05_Clip-5_REC709_2K_mask.mp4 b/example_workflows/PF0028-05_Clip-5_REC709_2K_mask.mp4 new file mode 100644 index 0000000..f039dc0 Binary files /dev/null and b/example_workflows/PF0028-05_Clip-5_REC709_2K_mask.mp4 differ diff --git a/example_workflows/corridorkey_example_workflow.json b/example_workflows/corridorkey_example_workflow.json new file mode 100644 index 0000000..d0988aa --- /dev/null +++ b/example_workflows/corridorkey_example_workflow.json @@ -0,0 +1 @@ +{"id":"93385f96-938f-405e-9213-1b6c5e29eddd","revision":0,"last_node_id":40,"last_link_id":49,"nodes":[{"id":31,"type":"ImageToMask","pos":[3590.5060043536255,-1027.9492878304789],"size":[270,58],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"IMAGE","link":35},{"localized_name":"channel","name":"channel","type":"COMBO","widget":{"name":"channel"},"link":null}],"outputs":[{"localized_name":"MASK","name":"MASK","type":"MASK","links":[36]}],"properties":{"cnr_id":"comfy-core","ver":"0.12.3","Node name for S&R":"ImageToMask"},"widgets_values":["red"]},{"id":30,"type":"GetVideoComponents","pos":[3326.5107507819284,-1024.8141943534986],"size":[193.530859375,66],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"video","name":"video","type":"VIDEO","link":34}],"outputs":[{"localized_name":"images","name":"images","type":"IMAGE","links":[35]},{"localized_name":"audio","name":"audio","type":"AUDIO","links":null},{"localized_name":"fps","name":"fps","type":"FLOAT","links":null}],"properties":{"cnr_id":"comfy-core","ver":"0.12.3","Node name for S&R":"GetVideoComponents"},"widgets_values":[]},{"id":35,"type":"MaskToImage","pos":[4332.245233792164,-1208.0833891802056],"size":[193.2779296875,26],"flags":{"collapsed":true},"order":7,"mode":0,"inputs":[{"localized_name":"mask","name":"mask","type":"MASK","link":41}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[42]}],"properties":{"cnr_id":"comfy-core","ver":"0.12.3","Node name for S&R":"MaskToImage"},"widgets_values":[]},{"id":34,"type":"CreateVideo","pos":[4544.55785001201,-1236.2810547128115],"size":[270,78],"flags":{},"order":11,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":42},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":null},{"localized_name":"fps","name":"fps","type":"FLOAT","widget":{"name":"fps"},"link":40}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[43]}],"properties":{"cnr_id":"comfy-core","ver":"0.12.3","Node name for S&R":"CreateVideo"},"widgets_values":[30]},{"id":37,"type":"CreateVideo","pos":[4525.482982815351,-1033.920507353679],"size":[270,78],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":44},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":null},{"localized_name":"fps","name":"fps","type":"FLOAT","widget":{"name":"fps"},"link":45}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[46]}],"properties":{"cnr_id":"comfy-core","ver":"0.12.3","Node name for S&R":"CreateVideo"},"widgets_values":[30]},{"id":27,"type":"GetVideoComponents","pos":[3296.7036105904185,-1594.7690495683291],"size":[193.530859375,66],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"video","name":"video","type":"VIDEO","link":32}],"outputs":[{"localized_name":"images","name":"images","type":"IMAGE","links":[33]},{"localized_name":"audio","name":"audio","type":"AUDIO","links":null},{"localized_name":"fps","name":"fps","type":"FLOAT","links":[39,40,45,47]}],"properties":{"cnr_id":"comfy-core","ver":"0.12.3","Node name for S&R":"GetVideoComponents"},"widgets_values":[]},{"id":39,"type":"CreateVideo","pos":[4515.530977228748,-874.6860135517358],"size":[270,78],"flags":{},"order":9,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":48},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":null},{"localized_name":"fps","name":"fps","type":"FLOAT","widget":{"name":"fps"},"link":47}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[49]}],"properties":{"cnr_id":"comfy-core","ver":"0.12.3","Node name for S&R":"CreateVideo"},"widgets_values":[30]},{"id":33,"type":"CreateVideo","pos":[4529.961473913082,-1411.4392161906396],"size":[270,78],"flags":{},"order":6,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":38},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":null},{"localized_name":"fps","name":"fps","type":"FLOAT","widget":{"name":"fps"},"link":39}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[37]}],"properties":{"cnr_id":"comfy-core","ver":"0.12.3","Node name for S&R":"CreateVideo"},"widgets_values":[30]},{"id":29,"type":"LoadVideo","pos":[2594.59614809964,-1015.7563952691187],"size":[441.9704841188509,436.37278717818185],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"file","name":"file","type":"COMBO","widget":{"name":"file"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[34]}],"properties":{"cnr_id":"comfy-core","ver":"0.12.3","Node name for S&R":"LoadVideo"},"widgets_values":["PF0028-05_Clip-5_REC709_2K_mask.mp4","image"]},{"id":26,"type":"LoadVideo","pos":[2597.563237437555,-1595.7749866942304],"size":[441.9704841188509,436.37278717818185],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"file","name":"file","type":"COMBO","widget":{"name":"file"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[32]}],"properties":{"cnr_id":"comfy-core","ver":"0.12.3","Node name for S&R":"LoadVideo"},"widgets_values":["PF0028-05_Clip-5_REC709_2K_input.mp4","image"]},{"id":28,"type":"CorridorKey","pos":[3953.987957010916,-1098.5768142940226],"size":[270,264],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"IMAGE","link":33},{"localized_name":"mask","name":"mask","type":"MASK","link":36},{"localized_name":"gamma_space","name":"gamma_space","type":"COMBO","widget":{"name":"gamma_space"},"link":null},{"localized_name":"despill_strength","name":"despill_strength","type":"FLOAT","widget":{"name":"despill_strength"},"link":null},{"localized_name":"refiner_strength","name":"refiner_strength","type":"FLOAT","widget":{"name":"refiner_strength"},"link":null},{"localized_name":"auto_despeckle","name":"auto_despeckle","type":"COMBO","widget":{"name":"auto_despeckle"},"link":null},{"localized_name":"despeckle_size","name":"despeckle_size","type":"INT","widget":{"name":"despeckle_size"},"link":null}],"outputs":[{"localized_name":"fg","name":"fg","type":"IMAGE","links":[38]},{"localized_name":"matte","name":"matte","type":"MASK","links":[41]},{"localized_name":"processed","name":"processed","type":"IMAGE","links":[44]},{"localized_name":"QC","name":"QC","type":"IMAGE","links":[48]}],"properties":{"aux_id":"SeanBRVFX/ComfyUI-CorridorKey","ver":"121bec58b0fb7b52eead028fbc606800382fa9bc","Node name for S&R":"CorridorKey"},"widgets_values":["sRGB",0.5,1,"On",400]},{"id":32,"type":"SaveVideo","pos":[4900.413131791065,-1669.330894601207],"size":[423.8060235788844,321.49145774667727],"flags":{},"order":10,"mode":0,"inputs":[{"localized_name":"video","name":"video","type":"VIDEO","link":37},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null},{"localized_name":"format","name":"format","type":"COMBO","widget":{"name":"format"},"link":null},{"localized_name":"codec","name":"codec","type":"COMBO","widget":{"name":"codec"},"link":null}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.12.3","Node name for S&R":"SaveVideo"},"widgets_values":["CorridorKey/FG/fg","auto","auto"]},{"id":40,"type":"SaveVideo","pos":[4911.008324925514,-481.61796156233123],"size":[439.1278487245063,329.57132647581386],"flags":{},"order":13,"mode":0,"inputs":[{"localized_name":"video","name":"video","type":"VIDEO","link":49},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null},{"localized_name":"format","name":"format","type":"COMBO","widget":{"name":"format"},"link":null},{"localized_name":"codec","name":"codec","type":"COMBO","widget":{"name":"codec"},"link":null}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.12.3","Node name for S&R":"SaveVideo"},"widgets_values":["CorridorKey/QC/QC","auto","auto"]},{"id":38,"type":"SaveVideo","pos":[4902.922731253981,-903.0076228787945],"size":[441.33536686144544,344.48230415353737],"flags":{},"order":12,"mode":0,"inputs":[{"localized_name":"video","name":"video","type":"VIDEO","link":46},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null},{"localized_name":"format","name":"format","type":"COMBO","widget":{"name":"format"},"link":null},{"localized_name":"codec","name":"codec","type":"COMBO","widget":{"name":"codec"},"link":null}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.12.3","Node name for S&R":"SaveVideo"},"widgets_values":["CorridorKey/Processed/Processed","auto","auto"]},{"id":36,"type":"SaveVideo","pos":[4903.709499337594,-1279.5522250408385],"size":[438.02375278465934,328.9890883825352],"flags":{},"order":14,"mode":0,"inputs":[{"localized_name":"video","name":"video","type":"VIDEO","link":43},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null},{"localized_name":"format","name":"format","type":"COMBO","widget":{"name":"format"},"link":null},{"localized_name":"codec","name":"codec","type":"COMBO","widget":{"name":"codec"},"link":null}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.12.3","Node name for S&R":"SaveVideo"},"widgets_values":["CorridorKey/Matte/matte","auto","auto"]}],"links":[[32,26,0,27,0,"VIDEO"],[33,27,0,28,0,"IMAGE"],[34,29,0,30,0,"VIDEO"],[35,30,0,31,0,"IMAGE"],[36,31,0,28,1,"MASK"],[37,33,0,32,0,"VIDEO"],[38,28,0,33,0,"IMAGE"],[39,27,2,33,2,"FLOAT"],[40,27,2,34,2,"FLOAT"],[41,28,1,35,0,"MASK"],[42,35,0,34,0,"IMAGE"],[43,34,0,36,0,"VIDEO"],[44,28,2,37,0,"IMAGE"],[45,27,2,37,2,"FLOAT"],[46,37,0,38,0,"VIDEO"],[47,27,2,39,2,"FLOAT"],[48,28,3,39,0,"IMAGE"],[49,39,0,40,0,"VIDEO"]],"groups":[],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":3.027374977196536,"offset":[-4668.655216692758,1528.3932379042944]}},"version":0.4} diff --git a/models/.gitkeep b/models/.gitkeep new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/models/.gitkeep @@ -0,0 +1 @@ + diff --git a/nodes.py b/nodes.py index 5b437ba..0527358 100644 --- a/nodes.py +++ b/nodes.py @@ -1,53 +1,82 @@ from __future__ import annotations +from typing import Callable + try: from .corridor_key import CorridorKeyProcessor, CorridorKeySettings except ImportError: from corridor_key import CorridorKeyProcessor, CorridorKeySettings +def _build_progress_reporter(unique_id: str | None) -> Callable[[str, int, int], None]: + progress_bar = None + prompt_server = None + + # Mini title: best-effort ComfyUI progress hooks + try: + from comfy.utils import ProgressBar + + progress_bar = ProgressBar(1, node_id=unique_id) + except Exception: + progress_bar = None + + if unique_id: + try: + from server import PromptServer + + prompt_server = PromptServer.instance + except Exception: + prompt_server = None + + def report(message: str, completed: int, total: int) -> None: + bounded_total = max(int(total), 1) + bounded_completed = max(0, min(int(completed), bounded_total)) + + if progress_bar is not None: + progress_bar.update_absolute(bounded_completed, total=bounded_total) + if prompt_server is not None and unique_id: + prompt_server.send_progress_text(message, unique_id) + + print(f"[CorridorKey] {message}") + + return report + + class CorridorKey: + DESCRIPTION = ( + "Refines a coarse per-frame alpha hint into CorridorKey FG, Matte, Processed, and QC passes. " + "The alpha hint should be rough, slightly eroded, and soft rather than tightly expanded." + ) + OUTPUT_TOOLTIPS = ( + "Raw straight foreground color. The model predicts this in sRGB space; convert to linear before manual compositing with the matte.", + "Raw linear alpha matte.", + "Linear foreground premultiplied by the linear matte for quick preview and simple downstream export.", + "QC preview composite over a checkerboard in sRGB.", + ) + @classmethod def INPUT_TYPES(cls) -> dict[str, dict[str, tuple]]: return { "required": { - "image": ("IMAGE",), - "mask": ("MASK",), - "gamma_space": (["sRGB", "Linear"],), - "feather_radius": ( - "INT", + "image": ( + "IMAGE", { - "default": 4, - "min": 0, - "max": 64, - "step": 1, + "tooltip": "Input RGB image or image batch to refine.", }, ), - "edge_focus": ( - "FLOAT", + "mask": ( + "MASK", { - "default": 1.5, - "min": 0.0, - "max": 4.0, - "step": 0.1, + "tooltip": ( + "Coarse Alpha Hint for the current frame. Keep it rough, soft, and slightly eroded. " + "For batched images, provide one matching mask per frame." + ), }, ), - "threshold": ( - "FLOAT", + "gamma_space": ( + ["sRGB", "Linear"], { - "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, + "tooltip": "Interpret the incoming image as sRGB or already-linear before the fixed 2048x2048 inference pass.", }, ), "despill_strength": ( @@ -57,33 +86,46 @@ class CorridorKey: "min": 0.0, "max": 1.0, "step": 0.01, + "tooltip": ( + "Green despill amount after inference. 0 disables despill, 1 is the standard default. " + "Higher values are stronger and can push heavy spill toward magenta or purple." + ), }, ), "refiner_strength": ( "FLOAT", { "default": 1.0, - "min": 0.1, + "min": 0.0, "max": 4.0, "step": 0.1, + "tooltip": "Scales the learned refiner delta. 1.0 is the standard default behavior.", + }, + ), + "auto_despeckle": ( + ["On", "Off"], + { + "tooltip": "Enable connected-component cleanup on the predicted alpha matte.", }, ), - "auto_despeckle": (["On", "Off"],), "despeckle_size": ( "INT", { "default": 400, - "min": 1, + "min": 0, "max": 4096, "step": 1, + "tooltip": "Minimum island area in pixels to preserve when auto-despeckle is enabled. 400 is the standard default.", }, ), - "output_mode": (["Processed", "FG", "Matte", "Comp"],), - } + }, + "hidden": { + "unique_id": "UNIQUE_ID", + }, } - RETURN_TYPES = ("MASK", "IMAGE", "IMAGE", "IMAGE", "IMAGE") - RETURN_NAMES = ("refined_mask", "selected_output", "foreground", "processed", "comp") + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE", "IMAGE") + RETURN_NAMES = ("fg", "matte", "processed", "QC") FUNCTION = "run" CATEGORY = "CorridorKey" @@ -95,26 +137,23 @@ class CorridorKey: 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, + unique_id: str | None = None, ): 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) + progress_callback = _build_progress_reporter(unique_id) + return self._processor.refine( + image=image, + mask=mask, + settings=settings, + progress_callback=progress_callback, + ) diff --git a/pyproject.toml b/pyproject.toml index be5cf05..dc8f3e3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,8 +11,11 @@ requires-python = ">=3.10" dependencies = [ # Keep runtime deps aligned with requirements.txt for ComfyUI installs. "numpy>=1.24", + "opencv-python>=4.8", "Pillow>=9.5", "torch>=2.0", + "torchvision>=0.15", + "timm>=1.0.24,<2.0", ] [project.optional-dependencies] diff --git a/requirements.txt b/requirements.txt index b0926c3..4bb429b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,5 +3,8 @@ # here matches common custom-node conventions and lets dependency installers # verify it is present without changing the node code. torch>=2.0 +torchvision>=0.15 +timm>=1.0.24,<2.0 numpy>=1.24 +opencv-python>=4.8 Pillow>=9.5 diff --git a/tests/test_processor.py b/tests/test_processor.py index be71806..3580e52 100644 --- a/tests/test_processor.py +++ b/tests/test_processor.py @@ -1,5 +1,6 @@ from __future__ import annotations +import numpy as np import pytest import torch @@ -9,21 +10,12 @@ from corridor_key import ( is_verified_check_conclusions, select_latest_verified_commit, ) +from corridor_key.color_utils import composite_straight, linear_to_srgb, srgb_to_linear +from corridor_key.tensor_ops import ensure_image_tensor, ensure_mask_batch +from nodes import CorridorKey -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: +def test_settings_match_upstream_defaults() -> None: settings = CorridorKeySettings() assert settings.gamma_space == "sRGB" @@ -31,135 +23,82 @@ def test_settings_match_upstream_like_defaults() -> None: 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_settings_reject_invalid_despill() -> None: + with pytest.raises(ValueError, match="despill_strength"): + CorridorKeySettings(despill_strength=1.5) -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_srgb_linear_round_trip_is_stable() -> None: + source = np.array([0.0, 0.018, 0.5, 1.0], dtype=np.float32) + round_trip = linear_to_srgb(srgb_to_linear(source)) + assert np.allclose(round_trip, source, atol=1e-5) -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_composite_straight_matches_expected_formula() -> None: + fg = torch.tensor([[[0.8, 0.2, 0.0]]], dtype=torch.float32) + bg = torch.tensor([[[0.2, 0.2, 0.2]]], dtype=torch.float32) + alpha = torch.tensor([[[0.5]]], dtype=torch.float32) + comp = composite_straight(fg, bg, alpha) + expected = torch.tensor([[[0.5, 0.2, 0.1]]], dtype=torch.float32) + assert torch.allclose(comp, expected) -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() +def test_ensure_image_tensor_requires_bhwc() -> None: 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) + assert ensure_image_tensor(image).shape == image.shape - _, _, 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]) + with pytest.raises(ValueError, match="shape"): + ensure_image_tensor(torch.zeros((1, 3, 4, 4), dtype=torch.float32)) -def test_auto_despeckle_can_remove_single_pixel_island() -> None: +def test_ensure_mask_batch_rejects_mismatched_batch() -> None: + mask = torch.ones((4, 4), dtype=torch.float32) + with pytest.raises(ValueError, match="exactly"): + ensure_mask_batch(mask, batch_size=2, height=4, width=4) + + +def test_node_metadata_populates_ui_info() -> None: + input_types = CorridorKey.INPUT_TYPES() + required_inputs = input_types["required"] + + assert "coarse" in required_inputs["mask"][1]["tooltip"].lower() + assert "per frame" in required_inputs["mask"][1]["tooltip"].lower() + assert input_types["hidden"]["unique_id"] == "UNIQUE_ID" + assert CorridorKey.RETURN_NAMES == ("fg", "matte", "processed", "QC") + assert len(CorridorKey.OUTPUT_TOOLTIPS) == 4 + assert CorridorKey.DESCRIPTION + + +def test_processor_reports_progress_messages(monkeypatch: pytest.MonkeyPatch) -> None: + class FakeEngine: + def process_frame(self, **_kwargs): + return { + "fg": np.zeros((4, 4, 3), dtype=np.float32), + "matte": np.zeros((4, 4, 1), dtype=np.float32), + "processed": np.zeros((4, 4, 3), dtype=np.float32), + "comp": np.zeros((4, 4, 3), dtype=np.float32), + } + + monkeypatch.setattr("corridor_key.processor.get_cached_engine", lambda device=None: FakeEngine()) + processor = CorridorKeyProcessor() - image = make_image() - mask = torch.zeros((1, 6, 6), dtype=torch.float32) - mask[:, 2, 2] = 1.0 + messages: list[tuple[str, int, int]] = [] + image = torch.zeros((2, 4, 4, 3), dtype=torch.float32) + mask = torch.zeros((2, 4, 4), dtype=torch.float32) - refined_mask, _, _, _, _ = processor.refine( + processor.refine( image=image, mask=mask, - settings=CorridorKeySettings( - auto_despeckle="On", - despeckle_size=9, - feather_radius=0, - preserve_core=0.0, - ), + settings=CorridorKeySettings(), + progress_callback=lambda message, completed, total: messages.append((message, completed, total)), ) - assert float(refined_mask[0, 2, 2]) == 0.0 + assert messages[0] == ("Loading CorridorKey model...", 0, 2) + assert messages[1] == ("Processing frame 1/2...", 0, 2) + assert messages[2] == ("Processing frame 2/2...", 1, 2) + assert messages[-1] == ("CorridorKey complete.", 2, 2) def test_verified_check_policy_requires_success_and_no_failure() -> None: