comfyui-dit-watermark

comfyui-dit-watermark

English | 简体中文

English

Add and detect invisible watermarks efficiently in ComfyUI image workflows. Detection requires neither diffusion inversion nor rerunning the generative model: a single VAE encode of the image is enough to recover the watermark. This makes the project fast and well suited to batch verification and automated pipelines. It provides GROW progressive frequency-domain watermarking nodes with flexible Reed–Solomon error correction, robust rotation/crop search, and RGB PSNR quality measurement.

| Source image | Identity clean | Watermarked | |---|---|---| | <img src="images/test/source_input.png" width="240" alt="Source image of a dog"> | <img src="images/demo/identity_clean.png" width="240" alt="Flux2 identity-clean image of a dog"> | <img src="images/demo/identity_watermarked.png" width="240" alt="Dog image with a GROW watermark"> |

The algorithm is based on luopengchen/GROW, with adaptations for ComfyUI's generic SAMPLER interface, Flux2 Klein 4B Distilled 4D latents, Qwen Image Edit single-frame 5D latents, and inference mode.

Three-image validation

The test uses the model, seed, sampler, frequency band, and prompt from the Flux2 UI workflow at workflows/flux2_klein_image_edit_grow.json:

Keep the input image exactly unchanged. Preserve every pixel, color, texture, composition, and detail.

The public input field is named watermark. All three images in this test contain zhangp36512345, with secret_key=watermark.

| Image | Clean → marked PSNR | ECC detection | |---|---:|---| | Dog | 41.212684 dB | Successful, 0 symbols corrected | | Claw | 40.970031 dB | Successful, 0 symbols corrected | | Girl | 38.193383 dB | Successful, 0 symbols corrected |

All three images fully recover zhangp36512345, and the incremental watermark PSNR is above the new 35 dB threshold in every case. The current unified workflow uses:

| strength | guidance_scale | start_ratio | Guided steps | |---:|---:|---:|---:| | 1.20 | 4000 | 0.00 | 4/4 |

PSNR uses a clean output produced with the same input, prompt, seed, model, and sampling parameters as the reference, so it measures only the change introduced by GROW. The source image also passes through Flux2/VAE reconstruction; differences between the source and output must not be counted as watermark distortion.

Nodes

GROW Watermark Config

Create one GROW Watermark Config, then connect its single GROW_CONFIG output to both the embedding and detection nodes:

GROW Watermark Config ─┬→ GROW DiT Sampler
                       └→ GROW Watermark Detect

This node stores all layout parameters that must match exactly, preventing subtle mismatches caused by entering the same values separately at both ends:

| Input | Purpose | Default | |---|---|---:| | secret_key | Determines the frequency-coordinate order | watermark | | dct_min, dct_max | Normalized mid-frequency range | 0.15, 0.45 | | max_channels | Number of latent channels used | 8 | | channel_start | First latent channel in the contiguous channel profile | 4 | | center_ratio | Central latent-area ratio used | 1.0 |

secret_key is stored as plain text in the workflow JSON. Protect workflow files in production, and never commit real secrets to a public repository.

GROW DiT Sampler

KSamplerSelect ─────────→ GROW DiT Sampler → SamplerCustomAdvanced
GROW Watermark Config ──→ GROW DiT Sampler

At selected denoising steps, the node intercepts the clean latent prediction x0 from the DiT. It computes an fp32 FFT/DCT-proxy sign-margin loss over secret-key-controlled mid-frequency coordinates, then returns the guided x0 to the original sampler. With the default start_ratio=0, guidance begins at the first step and runs for all 4/4 Flux2 sampling steps. The model, conditioning, scheduler, noise, and VAE weights remain unchanged.

| Input | Purpose | Flux2 default | |---|---|---:| | watermark | UTF-8 watermark content, up to 250 bytes | zhangp36512345 | | strength | Minimum signed frequency margin | 1.20 | | guidance_scale | Watermark gradient step size | 4000 | | start_ratio | Sampling ratio at which guidance starts | 0.00 | | config | Shared layout supplied by GROW Watermark Config | — |

The UI increments strength by 0.01 and guidance_scale by whole numbers. Watermarks longer than 32 UTF-8 bytes trigger a warning: longer frames reduce the number of frequency repetitions per bit, lowering attack robustness and increasing the number of candidate lengths during blind detection.

GROW Watermark Detect

Connect the generated IMAGE, the same VAE, and the same GROW_CONFIG used by the embedding node. The detector does not need to know the watermark content in advance because the flexible frame includes its length and Reed–Solomon validation data. max_watermark_bytes sets the blind-detection length limit, while robust_mode selects none, rotation, crop_scale, or a combined search.

Embedding and detection parameters must match

Together, secret_key, dct_min, dct_max, max_channels, channel_start, and center_ratio define how bits map to latent frequency coordinates. If any value differs at detection time, the detector reads different coefficients or splits bits differently and will usually fail RS validation. The detector does not need the watermark content, but it cannot arbitrarily infer the layout parameters from the current frame.

The current nodes require the sampler and detector to share these values through the same GROW_CONFIG, so the UI no longer exposes six separate layout fields at both ends. When detecting externally attacked images, you can perform a bounded search over a small set of predefined profiles (for example, 0–4), selecting the best result by RS validity and vote margin. Arbitrary brute-force searches over continuous dct_min/dct_max values are discouraged because the number of combinations, VAE runs, and false-positive surface all grow rapidly. To make an image fully self-describing, write a layout ID/header into a fixed pilot band that does not vary by profile, then use it to read the main payload.

A single-parameter mismatch test on Claw/profile 1 confirms this requirement. The exact match yielded zero raw codeword errors. Changing only channel_start to 0, max_channels to 4, dct_min to 0.10, or dct_max to 0.50 resulted in ecc_valid=False in all four cases. Note that the vote margin can still reach 1.0 when max_channels is mismatched, so a high margin does not prove that the layout is correct; RS/frame validation is authoritative.

Outputs include decoded_message, ecc_valid, corrected_symbols, min_vote_margin, and raw_codeword_hex. Detection follows this pipeline:

IMAGE → VAE encode → keyed latent frequency signs → majority vote → RS decode

It loads no additional diffusion model and performs no diffusion inversion.

GROW Image PSNR

Provide same-sized reference and watermarked IMAGE inputs to calculate RGB PSNR in the [0,1] range.

Reed–Solomon error correction

Each watermark is encoded as the shortest possible self-describing flexible frame:

1-byte UTF-8 length + N-byte payload + 4 RS parity bytes

The total length is N+5 bytes. The 14-byte watermark used here is reduced from the old protocol's 32 bytes to 19 bytes, giving each bit more frequency repetitions. Four parity symbols can correct any two corrupted byte symbols. If the damage exceeds this capability or the frame structure is invalid, the detector returns ecc_valid=False. Two ASCII characters correspond to two symbols; Chinese characters use multiple UTF-8 bytes, so correction capacity is measured in corrupted bytes. The default detection limit is 64 bytes, with backward compatibility for the old zero-padded 32-byte frame.

Attack resistance

• Flux2 Klein (Dog, Claw, and Girl): fully recovered 43 of 48 compressed, rotated, and cropped attack samples, for an 89.58% success rate. All AVIF, HEIF, JXL, crop, and JPEG q50-or-higher samples were recovered.
• Qwen Image Edit 2511 (initial Claw test): fully recovered 12 of 16 attack samples, for a 75.00% success rate. This is an initial single-image result and is not combined with the three-image Flux2 result.

See the attack-resistance benchmarks for complete test conditions, per-case results, resize comparisons, failure boundaries, and historical data.

Installation

Place the repository in:

ComfyUI/custom_nodes/comfyui-dit-watermark

Then restart ComfyUI. At runtime, the project depends only on the PyTorch bundled with ComfyUI; no additional ECC package is required. The remote deployment used for this validation was located at /app/custom_nodes/comfyui-dit-watermark.

Workflows

The repository includes two UI workflows that can be loaded directly into ComfyUI:

workflows/flux2_klein_image_edit_grow.json
workflows/image_qwen_image_edit_2511.json

• Flux2 Klein: flux2_klein_image_edit_grow.json targets four-step sampling with Flux2 Klein 4B Distilled.
• Qwen Image Edit 2511: image_qwen_image_edit_2511.json uses the fp8 model without loading a LoRA, and adds a connectable SAMPLER input to the original standard KSampler through advanced sampling components.

Both workflows include an identity prompt, a GROW Watermark Config connected to both embedding and detection, a GROW DiT Sampler, detection without prior knowledge of the watermark content, and watermark=zhangp36512345.

Compatibility and limitations

• Verified with ComfyUI 0.26.0, Flux2 Klein 4B Distilled, and four-step Euler sampling.
• The sampler wrapper can be reused with image DiTs that return 4D [B,C,H,W] latents or single-frame 5D [B,C,1,H,W] latents. Multi-frame video latents are not currently supported.
• The detector supports robust rotation and crop/scale search. Cropping is stable in testing, while rotation still produces noticeable residual errors.
• Compression resistance depends on image content; keep sufficient error-correction headroom for real-world images.
• secret_key is stored as plain text in workflow JSON files. Do not put production keys in public workflows.

Testing

python -m unittest discover -s tests -v

Tests cover shared GROW_CONFIG wiring and validation, GROW guidance, inference mode, flexible frames, legacy-frame compatibility, two-symbol correction, three-symbol rejection, blind length detection, geometric candidates, the public watermark interface, workflow configuration, and attack summaries.

Acknowledgements

Thanks to luopengchen/GROW for open-sourcing the progressive frequency-domain watermarking algorithm and implementation on which these ComfyUI nodes are based.