Envy Anti-Drift

ComfyUI Envy Anti-Drift

Envy KSampler Advanced Anti-Drift restores small blurred RGB differences to the original image while retaining larger changes. It adds sparse correction steps, mask expansion, optional mask retention, and a final image composite to the native KSampler Advanced controls.

This pack also includes Text Encode Qwen Image 2.1 (Experimental), moved from ComfyUI-Qwen21-Experimental with the same node ID and behavior. See its encoder documentation.

Clone ComfyUI-Envy-Anti-Drift into ComfyUI/custom_nodes and restart ComfyUI. No additional dependencies are required.

Connect the original image being edited to image1 and the matching VAE to vae. If image1 contains a batch, its first image is the reference for every sampled latent. When pixel dimensions differ, the reference is aspect-preserving center-cropped and resized with Lanczos to the decoded prediction; matching dimensions avoid resampling. Alpha is ignored.

At each selected correction, the node decodes the clean prediction after CFG. It compares the mean absolute RGB difference between independently filtered copies of the prediction and a VAE-reconstructed reference. The reference is encoded and decoded once per run at the first known prediction size, including runs without intermediate corrections, to reduce masks caused by VAE reconstruction differences. Pixels above threshold_percent seed the binary edit mask; optional connected growth adds weaker nearby differences. Small binary blobs are removed before mask_expand_pixels expands the mask and blur_pixels feathers it. ECC estimates translation and independent horizontal and vertical scale from the raw prediction against the original reference. The node aligns the entire raw prediction, including edited pixels, then applies the same background-based percentile gain/offset used by the final pass and blends the aligned prediction inside with the original reference outside the mask. The corrected composite is VAE-encoded and returned to the sampler as the corrected clean prediction. Previews and intermediate frames include this color correction.

comparison_filter defaults to gaussian. bilateral smooths similar colors while preserving stronger color edges, so fine detail added or removed along a line can remain detectable. comparison_blur_pixels is the comparison filter's spatial standard deviation (default 2 pixels); 0 compares unfiltered images. comparison_color_tolerance_percent is the bilateral RGB color-range standard deviation (default 10%), not a hard cutoff. Larger values allow more smoothing across color differences; this is separate from threshold_percent, which defaults to 5% of the normalized RGB range and sets the edit threshold. Bilateral comparison runs in halo tiles to reduce intermediate memory. CPU images use ComfyUI's configured CUDA or ROCm device for each tile when available; CPU-only mode stays on CPU. The filtered image retains the input device and dtype. Tile size targets at most 256 MiB of estimated temporary work and one-quarter of reported free accelerator memory, excluding CUDA context and allocator overhead. Very wide blur still needs memory proportional to its kernel. blur_pixels controls Gaussian mask feathering only (default 6 pixels); 0 disables feathering. The original reference, its VAE reconstruction, and the filtered reconstruction are cached between intermediate corrections.

mask_expand_pixels defaults to 15 pixels in every direction using a square footprint, limited to the existing image bounds. 0 disables expansion. Expansion occurs before feathering.

min_mask_area defaults to 8 decoded-image pixels. Before expansion, each batch item's binary mask drops blobs with fewer pixels. Only horizontal and vertical neighbors connect; diagonal neighbors do not. Blobs of eight pixels or more survive at the default. 0 or 1 disables this cleanup. Reused and retained masks are cleaned before display and storage. The provisional exclusion mask used to estimate ECC alignment is unchanged.

mask_growth defaults to true. threshold_percent selects strong seed pixels; strong blobs smaller than min_seed_area (default 128 pixels) are discarded. The mask then grows through pixels above growth_threshold_percent (default 2%) that connect to a surviving seed by horizontal or vertical steps and lie within max_growth_distance (default 64 straight-line pixels) of a seed. Strong seeds survive even if the growth threshold is higher. Zero growth distance keeps only seeds; zero or one seed area disables seed cleanup. Growth happens before retained masks are combined and before min_mask_area, expansion, and feathering. ECC's current provisional detector still uses the strong threshold. A weak line separated from every seed by a below-threshold or diagonal-only gap will not grow. Turn mask_growth off for the original single-threshold mask.

ECC aligns the unmasked image background with independent X/Y translation and scale. Pixels above the current difference threshold and any saved unexpanded binary edit mask are excluded from registration. Expansion and feathering affect compositing only, not registration exclusion. If too little reliable background remains, the raw prediction stays in place. When the maximum displacement across the image corners exceeds one-third of a pixel, the node recalculates the edit mask from the aligned raw prediction, excluding newly exposed edge pixels. Exposed edge pixels are filled from the original image after matching its background colors to the aligned prediction, with a two-pixel feather into valid pixels. The mask stays in output coordinates and is zero on exposed pixels. The original background remains at its original coordinates during compositing. Retained binary masks stay in output coordinates for later corrections. The final decode gets its own transform estimate and recalculates its mask above the same corner-displacement threshold; don't unmask pixels controls whether it retains earlier mask locations.

correction_steps defaults to 0.4,0.999. Separate values with spaces or commas; mixed inputs such as 4, 0.8 work. Whole numbers are zero-based original KSampler Advanced step positions. Fractions strictly between 0 and 1 multiply the full configured steps and round down: with 25 steps, 0.999 selects step 24, and 0.999, 24 corrects only once. Duplicates are removed after conversion. Positions outside the actual sampled range are ignored, including positions skipped by start_at_step. An empty string disables intermediate corrections; invalid values raise an error.

Each selected position runs once, at the first evaluation in its sigma interval. Predictor/corrector samplers such as Heun can enter that interval during the preceding outer step; repeated evaluations in the same interval do not trigger another correction. Other evaluations pass through without VAE work, history frames, or mask updates. Sampling remains one native call, preserving solver history and all native noise, start/end, and leftover-noise controls.

don't unmask pixels defaults to false. When enabled, each batch item's binary edit mask accumulates all locations detected at selected corrections that survive blob cleanup. The retained binary mask stays in the output image's coordinates and is unexpanded and unfeathered; expansion and blur are applied afresh, avoiding repeated outward growth. Only mask locations persist, while predicted pixels keep evolving. Early false positives that survive cleanup remain selected. With retention disabled, the final pass reuses the last mask unless corner displacement exceeds one-third of a pixel.

During correction, live previews reuse the corrected composite at full decoded resolution, using the first image of the sampled batch. Other evaluations use normal configured latent previews. The final composite also receives a full-resolution preview. Previews have no magenta overlay; masks remain grayscale in the second output.

The first output is the native sampled LATENT. intermediate_steps_and_masks is a diagnostic IMAGE batch with three side-by-side pairs for each executed correction and the final decode, even if no intermediate correction ran. Frames are ordered by correction evaluation (then Final), batch position, and pair: raw VAE decode | aligned raw decode; corrected composite | actual displayed feathered mask; registration exclusion mask | authoritative unexpanded, unfeathered binary edit mask. White in the feathered mask keeps prediction; white in the binary mask marks detected or retained edits. White in the registration exclusion mask marks pixels omitted when estimating the transform, including the previous binary mask. The aligned raw panel shows uncovered edges as black; mask recalculation fills those edges from the reference and excludes them from detection.

Each frame has a 32-pixel label band above the unchanged full-resolution image pixels. The label identifies the correction number and sigma (or Final), batch position, sampling scales sx/sy and offsets dx/dy, whether the mask was recalculated from aligned pixels, and whether the binary mask was detected, redetected, or reused. A positive sampling offset reads source pixels from a larger coordinate, so the displayed content moves in the opposite direction. Diagnostic frames are CPU float32 and require about 72 bytes per decoded pixel per correction and batch item, plus the final pass and a copy when assembling the output batch. The labels and aligned raw view use CPU image operations; they add no VAE calls and do not alter the live preview.

The third output, final_image, decodes the returned latent once, aligns it, and composites it with the original image in CPU float32. It reuses the last binary mask when maximum corner displacement is at most one-third of a pixel; above that, it detects a fresh mask against the VAE-reconstructed reference, combining it with the saved mask when don't unmask pixels is enabled. If no intermediate correction ran, it detects a mask from this final decode using the same reconstructed reference. Expansion and feathering are applied afresh. Before restoring the source background, it compares the aligned prediction and original source pixels where the feathered mask is exactly zero. Per-image, per-channel 10th/50th/90th percentiles define a color gain and offset, which are applied to the aligned prediction. Empty background samples or zero percentile spread in either channel leave that channel unchanged.

The composite copies the original pixels exactly wherever the feathered mask is zero. Feathered edges mix source and aligned prediction pixels. The original background is never transformed; the exposed edge fill is matched to the aligned prediction. The final image is never re-encoded, and the first output remains the native sampler result. VAE reconstruction can still alter the LATENT outside the mask. When the source needs resizing, exact preservation refers to its Lanczos-aligned pixels. Color matching may produce values outside 0–1.

raw_final_image is the aligned, background-color-matched final prediction immediately before blending with the original; it may differ from the original outside the edit. final_mask is the exact effective feathered blend alpha as a single-channel MASK batch: white keeps raw_final_image, black keeps the original. These fourth and fifth outputs reuse the final decode and color match, with no additional VAE work.

Each run adds one VAE encode and decode for the comparison reference. Each intermediate correction adds one VAE decode and encode, and finalization adds one decode. Model transfers between RAM and VRAM can dominate the cost. The final image is composited even when intermediate corrections are disabled or no selected position was sampled.

The threshold measures change size, not edit intent: it can preserve large unwanted changes or erase subtle intended edits. Early predictions can differ across most of the image and therefore receive little correction.

Requires ComfyUI 0.40.0 or newer and a compatible image VAE with RGB image decoding. Tested with ComfyUI 0.40.0 using its native model-patcher, sampler-wrapper, and memory-management APIs; it makes no changes to ComfyUI core.

ECC maximizes normalized intensity correlation with four parameters (X/Y translation and independent X/Y scale), using SciPy already required by ComfyUI. It does not estimate rotation, shear, or local deformation.

Licensed under GPL-3.0-or-later. This node builds on ComfyUI's native KSampler Advanced; see LICENSE.

Tests

From the pack directory, run the CPU tests with ComfyUI available on PYTHONPATH:

PYTHONPATH=/path/to/ComfyUI python -m pytest --import-mode=importlib -q tests/test_anti_drift_sampler.py