Nodes/comfyui-dit-watermark/GROW DiT Sampler
ComfyUI Node

GROW DiT Sampler

Bakes a recoverable watermark into the image while it generates

By zhangp365·Created 2 months ago·Updated 2 months ago· 1
GROW DiT Sampler
  • sampler
  • config
  • sampler
◄watermarkzhangp36512345►
◄strength1.20►
◄guidance_scale4000►
◄start_ratio0.00►

Most "invisible watermark" tools fail because they stamp a pattern over the finished image. Crop it, recompress it, resize it, and the watermark dies with the pixels. GROW DiT Sampler takes the opposite route: it bakes the payload into the image during denoising, at the level of the model's latent representation, so it rides along far better through the attacks that kill overlays.

It's the heart of the comfyui-dit-watermark pack - a ComfyUI port of luopengchen/GROW's progressive frequency-domain watermarking. This node wraps any existing SAMPLER and adds GROW guidance at chosen denoising steps. The author validates it on Flux2 Klein 4B Distilled and Qwen Image Edit, but it's a generic wrapper: anything that produces 4D [B,C,H,W] latents (or a single-frame 5D [B,C,1,H,W]) works. Multi-frame video latents don't, and the node will tell you so.

What it actually does

Each guided step, the node intercepts the model's clean latent prediction - the x0 the sampler is about to refine - and runs one gradient step of a sign-margin loss in the frequency domain. It uses the real part of an orthonormal FFT as a differentiable DCT stand-in, which is the classic GROW trick. Your secret_key seeds a scrambled order of mid-frequency coordinates, so each bit of the payload gets mapped onto several frequency coefficients; the loss pushes weak or wrong-signed coefficients toward the correct sign and leaves already-correct ones untouched. Then it hands the guided latent back to the original sampler. Model weights, conditioning, and the noise schedule are never touched - that's why the image stays otherwise identical.

With the default start_ratio=0, guidance runs from the first step through all 4/4 Flux2 steps. Because the whole thing is one gradient step per step, it costs a few extra autograd passes, not a rerun of the model.

The inputs that matter

Only three are worth touching on day one:

  • watermark - the payload, up to 250 UTF-8 bytes (default zhangp36512345). Past 32 bytes you get a warning: longer payloads mean fewer frequency repetitions per bit, which hurts attack robustness and slows blind detection.
  • strength - the minimum signed frequency margin, default 1.20. The UI steps it by 0.01.
  • guidance_scale - the gradient step size, default 4000. The UI steps it by whole numbers.

start_ratio (0–0.99) decides when guidance begins - 0 means "first step", which is what the validation workflows use. These defaults are tuned for Flux2 Klein's 4-step distilled schedule; on a different model you're on your own, but the two knobs above are the obvious thing to raise if the detector comes back empty.

Wiring it up

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

It's a transparent wrapper, so it slots between KSamplerSelect and SamplerCustomAdvanced with zero changes to the rest of the graph. The single output is another SAMPLER.

Installing

The whole pack installs at once (this is the only package with all four nodes):

cd ComfyUI/custom_nodes
git clone https://github.com/zhangp365/comfyui-dit-watermark

Then restart ComfyUI. ComfyUI Manager users can just search comfyui-dit-watermark. There are no model downloads and no extra pip packages - the pack's requirements.txt is deliberately empty because it only needs the PyTorch ComfyUI already bundles, and it rolls its own Reed–Solomon codec. The README's validation ran on ComfyUI 0.26.0.

Where people get burned

  • Sampler and detector config must match. secret_key, the frequency band, channels, everything comes from one GROW Watermark Config - wire the same config output into both nodes, or detection quietly reads the wrong coefficients.
  • Tuning a new model. The defaults are Flux2-Klein-tuned. If ecc_valid stays False, try raising strength first, then guidance_scale.
  • Your secret lives in plaintext. The workflow JSON - and any PNG you share, since ComfyUI embeds the workflow in the file - contains secret_key verbatim. Keep real keys out of public workflows.
CategoryGROW/watermark

Inputs (6)

NameTypeDefaultDescription
samplerSAMPLER—
configGROW_CONFIG—
watermarkSTRINGzhangp36512345—
strengthFLOAT1.200.01–5—
guidance_scaleFLOAT40001–20000—
start_ratioFLOAT0.000–0.99—

Outputs (1)

NameTypeDescription
samplerSAMPLER—