TBG Flux2 Differential Diffusion Inpainting

Flux 2 is a 4MP model with a proper edit mode, but standard inpaint masks don't play nice with it - mask-based correction applied too early in sampling fights the model's distilled sweet spot and you get mud. This node is TBG's fix: a hook you place before any sampler. It stores your inpaint mask privately and wraps the sampler so that mask correction only kicks in at the sigma where Flux 2 actually benefits from it. The node description sums it up: "Place this before any sampler... wraps a SAMPLER so Euler, LCM, DPM++, RES/other custom samplers can delay mask correction until Flux2's sweet spot."

It comes from the TBG ETUR pack (Tobias Laarmann / TBG), which has been shipping Flux 2 support since the model launched - the author's own Reddit post about the Flux 2 branch is one of the earliest community workflows for it, and it warns you'll want a big card (24–32 GB VRAM territory).

How it works

Under the hood it's built on ComfyUI's KSamplerX0Inpaint, which implements differential diffusion: instead of treating the mask as a hard "paint inside this shape" boundary, the mask's gray values control how much each region is edited, transitioning progressively across sampler steps. The node's mask tooltip says exactly that: white edits, black preserves, gray transitions.

The twist is timing. Flux 2 is a distilled model, and the pack's sampler code is built around a fixed distilled timestep ladder - the same 0.999/0.866/0.634/0.342/0.0 schedule the PiD samplers use. There's a "sweet spot" sigma (default correction_start_sigma = 0.6341) where mask correction actually helps. The hook waits until both the sampler has reached that sigma and the differential denoise timing allows it, then starts correcting. Apply the mask earlier and you're fighting the model's distillation; apply it later and edits look half-baked.

Inputs

• model and latent - your Flux 2 model and the latent you're editing. The hook passes both through, patched.
• inpaint_mask - the mask. Remember: white = edit, black = preserve, gray = gradual transition. This is the node's whole reason to exist.
• denoise (1.0) - this is differential inpaint timing for the hook itself, not the downstream sampler's denoise. The tooltip spells it out: the downstream sampler can run lower denoise; this value delays preservation inside the wrapped sampler.
• correction_start_sigma (0.6341) - the Flux 2 sweet spot. If edits look under-corrected, nudge it earlier; if they look over-smudged, push it later.
• sampler (optional, recommended) - connect any sampler here (Euler, DPM++, LCM, a RES4LYF sampler, whatever) and use the wrapped sampler output. Without it, the hook still patches model/latent but you lose the full sweet-spot latent correction.

Outputs: model, latent, and sampler - all three flow downstream into your KSampler as usual. The wrapped sampler is the one that carries the timing logic.

Install

Part of the TBG pack:

cd ComfyUI/custom_nodes
git clone https://github.com/Ltamann/ComfyUI-TBG-ETUR
cd ComfyUI-TBG-ETUR && pip install -r requirements.txt

or ComfyUI Manager → "TBG Enhanced Upscaler," restart. Heavy requirements (kornia, timm, triton, transformers/diffusers for PiD) and a current ComfyUI with the v3 API are required. You also need Flux 2 weights yourself - the README lists flux2 files, and per the author, Flux 2 runs want serious VRAM.

Common issues

• Edits look weak - you probably haven't connected the sampler input, so the wrapped sampler never runs. It's optional but "recommended," and it's where the actual timing logic lives.
• "Why is my mask doing nothing in the first steps?" - that's the feature. Mask correction is deliberately delayed to Flux 2's sweet spot. If it feels too slow to act, lower correction_start_sigma.
• Not a Flux 2 user? - this hook is tuned to Flux 2's distilled schedule. Don't expect it to be magic on a plain Flux 1 or SDXL inpaint; that's not the target.