AIMZ VAE Color Match

You run a video through an upscale pass, decode it, and the colors have drifted - a purple shift, a washed-out look, contrast in the wrong place. It's a real and common failure. The community talks about "color drift" constantly when chaining i2v segments or upscaling video: people report exactly this purple shift when they chain WAN2.2 segments, and the fix is never "re-render." AIMZ VAE Color Match is a targeted answer: feed it the original frame as a color reference, feed it the drifted result, and it transfers the original's color statistics back onto the new image.

This is the one node in the pack with actual signal-processing weight behind it, and it's aimed squarely at LTX 2.5 upscale pipelines - the README's own framing is "1st pass" reference vs. the "VAEDecode (LTX 2.5)" output that comes out with drifted color.

How it works

Frame by frame, on the GPU, it computes a color transform from the reference image and applies it to the processed image. There are five algorithms, and the names tell you what's happening:

• statistical_matching (GPU) - the default. Per-channel RGB moment transfer: shift the mean, scale the standard deviation, so the processed frame's channel stats match the reference. Fast and clean.
• exact_histogram (GPU Powerful) - sorts both images' per-channel values and remaps the processed pixels onto the reference's cumulative histogram. Exact, but harder-edged.
• lab_reinhard (GPU Perceptual) - the same mean/std transfer, but in LAB color space, which is closer to how human vision perceives color and tone.
• lab_histogram (GPU Cinematic) - histogram transfer in LAB, described as a "cinematic" grade.
• luminance_zones (GPU) - just the brightness-bias correction, for when only exposure drifted.

A correction_strength slider blends between the untouched processed image (0.0) and the full match (1.0, the default). If the reference and processed frames differ in resolution, the reference is resized to match. If your reference video is shorter than the processed one, it loops the reference - the node maps each processed frame to min(i, ref_frames-1).

Inputs and outputs

The ones that matter: original_image (your reference, e.g. the loaded source video or 1st-pass output), processed_image (the drifted result), correction_strength, method, and auto_preserve (keeps heavily-altered regions - roughly the top quarter of per-pixel difference - untouched, so the correction doesn't stomp areas the model legitimately changed). Optional mask lets you preserve areas manually, and edge_feather is meant to feather the boundary between corrected and preserved areas.

The single output is corrected_image, ready to feed your saver or the next stage.

Installing it

Part of the ComfyUI_AIMZ pack. ComfyUI Manager → Custom Nodes Manager → search comfyui_aimz → Install → restart. Or:

cd ComfyUI/custom_nodes
git clone https://github.com/AIMZ-GFX/ComfyUI_AIMZ.git

No model downloads. Everything is torch/numpy math, GPU-side.

Honest caveats

Two things from the source that the README glosses over. The optional vae input is accepted but not actually used in the computation - it's a convenience for keeping your pipeline wired, not a dependency. And edge_feather is declared as a widget, but the shipped code (v0.5.7) doesn't apply it yet, so don't expect soft transitions from it until a later release.

The bigger practical caveat: this matches color, not content. If the upscale genuinely changed the image, the color match makes the new pixels wear the old colors - which is usually exactly what you want, but don't expect it to restore structure. And the mask note from the code: white areas in the mask are where the correction lands, black keeps the processed input as-is. If your mask is inverted, the result will look backwards.

Start with the default statistical match and strength at 1.0; if skin tones look off, that's when you drop to lab_reinhard, which is gentler on how colors are perceived. It's a small pack, but this node is the reason people install it.