Image Color Restore
Put the original palette back after an AI edit
- edit_image
- org_image
- image
- unchanged_mask
The problem it solves
You inpaint a jacket, or drop an object out, or run the whole frame through an instruction editor for one small change. The edit itself looks right - and the image comes back slightly off. A haze, a warmth shift, a bit less contrast than the reference. It's the complaint that drove the whole crop-and-stitch workflow around editing models: edit models re-emit the entire frame, so pixels nobody asked about move.
Image Color Restore is the fix for the colour half of that drift, specifically for the case where you have two images: the AI edit, and the aligned original you edited from. It fits a colour transform mapping one onto the other, applies it, and locally smooths the seams left behind where objects were removed.
Colour matching is the underrated primitive of the post-processing layer - deterministic, instant, and the honest alternative to re-rolling a generation because "the colours are off". This node is the aggressive, reference-aware version of that idea; for a plain global statistics transfer, the same pack's Image Color Match does that (it wraps the color-matcher library).
How it works
Four stages, and the first is the one that keeps it from wrecking your actual edits.
Foreground/background split. The edit is segmented in LAB space against its background. The author is blunt about this being a hard requirement: edit_image must place the subject on an approximately solid background, or segmentation fails and the whole fit fails with it. If the foreground comes out below 0.5% or above 80% of the frame, the node raises rather than returning something subtly wrong.
A trusted affine colour fit. A transform is solved from edit_image to org_image using only pixels confidently unchanged between the two. It's a robust (Tukey IRLS) fit that iteratively zeroes out samples with big residuals, so the region you actually edited is excluded instead of being averaged into the model and dragging the whole correction toward it.
A continuous seam field. Residual differences are propagated outward from the trusted unchanged region, with a distance falloff and a Lab colour-similarity gate, so the compensation fades in smoothly rather than stepping at the mask boundary.
An occlusion pass. Where a removed object is detected, a wider residual blur is applied locally; set occlusion_residual_blur or occlusion_bridge_width to 0 to disable that pass.
Credit where it's due: the algorithm is a port of ComfyUI-RH-Nodes' Reference Color Restore (Occlusion Seam Advanced) V0.8, inlined into this pack rather than required as a dependency.
Inputs and outputs
edit_image and org_image are the two that matter, and both must be pixel-aligned. If the sizes differ, edit_image is resized down/up to org_image - a few pixels of drift is tolerated, a different framing is not.
background_threshold (0–80, default 0 = auto) - the LAB distance used for the foreground/background split. Leave it on auto unless segmentation is failing on a soft or busy background; this is the knob to reach for when the node errors out about foreground fraction.
unchanged_threshold (0–80, default 0 = auto) - the maximum RGB residual for a pixel to still count as "unchanged". Higher means more of the image counts as unchanged, which restores the reference more aggressively. This is the one an advanced user actually tunes.
The rest - seam_max_distance, seam_decay, seam_residual_blur, seam_color_sigma, occlusion_residual_blur, occlusion_bridge_width - control how far the seam compensation reaches and how it's blurred. The defaults (96/64/24/20/48/48) are sensible; touch them if you see a visible halo around the edit.
Outputs: image, the colour-restored result, preserving edit_image's alpha channel when it had one; and unchanged_mask, a MASK marking the pixels the edit left alone - precisely the pixels the output took from org_image rather than from the colour-mapped edit.
unchanged_mask is more useful than it sounds. Feed it to a blend or paste node and you can decide separately what happens in the untouched regions versus the edited ones, instead of trusting the node's own composite.
Install
Ships in ComfyUI 1hewNodes. Manager → search the pack title, or:
cd ComfyUI/custom_nodes
git clone https://github.com/1hew/ComfyUI-1hewNodes
pip install -r ComfyUI-1hewNodes/requirements.txt
Restart ComfyUI. The algorithm is pure NumPy/OpenCV, so this node needs no model downloads - the long requirements list belongs to the pack's background-removal, PSB and video nodes.
Common issues
"Foreground must end up 0.5%–80%." The most common failure, and it's the input's fault, not the node's. Either your subject fills nearly the whole frame, or the background isn't solid enough for the LAB split. Crop some margin around the subject, or nudge background_threshold off auto.
Misalignment beats everything. A colour fit can survive a global exposure shift; it cannot survive the edit model re-framing the subject. Run the alignment step first - the same pack's Image Align Change Mask is built for exactly this pairing - then restore colour.
Don't run it on a whole-frame restyle. This node assumes "local edits plus slight global drift." If the edit changed the scene's colour on purpose, the node will faithfully undo your intent.
Inputs (10)
| Name | Type | Default | Description |
|---|---|---|---|
| edit_image | IMAGE | AI-edited image on a solid background; must be pixel-aligned with org_image (resized to org_image's size when they differ) | |
| org_image | IMAGE | — | |
| background_threshold | FLOAT | 0.00–80 | LAB distance threshold for foreground/background split; 0 = auto. Foreground must end up 0.5%-80% or the node errors. |
| unchanged_threshold | FLOAT | 0.00–80 | Max RGB residual (0-255) for a pixel to count as 'unchanged'; 0 = auto. Higher restores the reference more aggressively. |
| seam_max_distance | FLOAT | 961–1024 | Max propagation distance (px) from trusted unchanged region |
| seam_decay | FLOAT | 641–1024 | Distance falloff exp(-d/decay); larger reaches farther |
| seam_residual_blur | FLOAT | 24.00.1–128 | Gaussian sigma for smoothing the residual field |
| seam_color_sigma | FLOAT | 20.00.1–128 | Lab color-similarity gate sigma (only propagates to similar colors) |
| occlusion_residual_blur | FLOAT | 48.00–128 | Extra residual blur used only in the detected occlusion zone (0 disables local pass) |
| occlusion_bridge_width | FLOAT | 48.00–128 | Occlusion detection reach / bridge width (0 disables local pass) |
Outputs (2)
| Name | Type | Description |
|---|---|---|
| image | IMAGE | — |
| unchanged_mask | MASK | — |