VAE Clean — cast / saturation / decode artifacts

The complaint this node was built for

"The colours came back reddish, oversaturated and too contrasty, and it is not the prompt." That is usually true, and it is usually four different faults wearing one coat. A VAE with a constant per-channel bias reads as a cast and divides out exactly. Chroma pushed too far reads as oversaturation and is very often CFG rather than the decoder. A crushed black point and an over-extended white point read as excess contrast and are only recoverable where they have not clipped. And clipping itself cannot be undone at all - a channel pinned at 1.0 has lost the values above it, and any tool that "recovers" it is inventing detail.

Most cleanup nodes jump straight to correcting. This one measures first, tells you which of those it found, then applies only what you asked for.

How it works, and why linear matters

Everything happens in linear light: real piecewise sRGB decode, operate, re-encode. That is not pedantry. Pull 20% of the saturation out of sRGB values and you are applying the operation to a display curve, which is why a naive saturation slider always seems to move exposure as well. Doing it about Rec.709 luma instead keeps brightness where it was.

balance_mode is the correction half. off measures only and changes nothing - start there and read the report. grey world drives the average neutral: fast, and wrong on any shot that is genuinely one colour, because a sunset averages red and this will throw the sunset away. white point drives the brightest 1% neutral, which is safer on coloured scenes since a red room still has a neutral highlight. match reference copies a plate's channel means and needs the optional reference input wired. balance_strength (0–1) dials the whole thing back when the cast is partly the lighting you wanted.

saturation (0–2) scales chroma about luma. contrast_restore (0–1) lifts crushed shadows - which reduces contrast, the fix for a render that came back too hard. It raises the floor rather than stretching the range, so nothing can clip, and it does not resurrect detail already crushed to black; the report tells you how much is in that state. chroma_cleanup (0–1) blurs chroma only, which is the fix for VAE decode speckle and can go surprisingly hard before anything looks soft, because the eye carries detail in luma.

restore_unchanged (0–0.5, experimental) is the write-back: with the original plate on reference, it restores the reference's pixels wherever the picture did not really change and keeps your edit. The value is how far the local average may move and still count as unchanged - on the Flux VAE, one round trip moves it up to 0.037 and three up to 0.050, so start near 0.05–0.07.

Outputs

image (corrected), report (a text diagnosis - read it), cast_strength and saturation (floats, if you want to drive something else off the measurement).

Install

cd ComfyUI/custom_nodes
git clone https://github.com/Code2Collapse/ComfyUI-CustomNodePacks.git

ComfyUI Manager → "CustomNodePacks" does the same. Restart and confirm [C2C] CustomNodePacks: 142 nodes loaded (...) - 0 failed.

No models, no downloads, nothing GPU-specific - this is arithmetic on tensors. The pack still wants its shared packages, installed carefully so you do not overwrite ComfyUI's torch:

pip list | grep -i "opencv\|scipy\|safetensors"
pip install opencv-python>=4.7.0 scipy>=1.10.0 safetensors>=0.4.0

Where people get burned

• Reaching for this when the real cause is CFG. If a render is oversaturated and contrasty, that combination is very often guidance, and the author says so in the tooltip. On 2026's guidance-distilled models, CFG 1 is the intended default and raising it applies the correction twice, burning the image; desaturating afterwards cannot bring that detail back. Set CFG to 1 (or 1.2–1.5 with more steps if anatomy is failing) and re-render first.
• grey world on a sunset. It will neutralise the shot, beautifully and wrongly. Use white point or match reference.
• Expecting clipping recovery. It reports the clipped share instead of pretending. That is the correct behaviour and it means the answer is "re-render with less CFG", not "clean harder".
• VAE loss is not a colour problem. Round trips are lossy - every encode/decode degrades the image a little, and chained passes accumulate it. If that is your issue, restore_unchanged with the plate wired to reference is the tool, and so is not decoding and re-encoding in the first place.
• Installing it expecting a quality upgrade. The pack's own VAE docs are clear that fp32 is not where visible VAE loss comes from; this node is a diagnostician for colour shifts, not a magic decoder. It sits next to the pack's VAE Quality Decode node, not on top of it.