Flux2 JiT Sampler

Flux2JiTSampler has exactly zero inputs. Double-click it and there's nothing to configure - just a SAMPLER output. If you're new to ComfyUI's custom-sampling path, a node with no knobs looks like a stub. It's not. This is where JiT's stage transitions actually happen, and it's the node that refuses to let you run the pipeline broken.

Why it exists

JiT isn't just "evaluate fewer tokens per step." The tricky part is the stage transitions: when more tokens come online mid-run, they have to be grafted onto the real evolving latent state with the correct noise level - the paper calls it the deterministic micro-flow (DMF) target. A plain model wrapper can sparsify the model call, but it can't touch the sampler's latent. So the pack splits the work: Flux2JiTApply patches the model with the sparse wrapper, and this sampler owns the transitions, reusing the same noise across them so newly activated tokens don't enter with fresh random noise (that was a bug in the first port, now fixed). They're inseparable - that's why the README wires the Apply output and this sampler into the same SamplerCustom.

How it works

Under the hood it's a comfy.samplers.Sampler doing Euler-style one-eval-per-step sampling - the intended path for the whole pack. It also self-checks: if JiT stage logic is active but the sparse wrapper didn't actually execute a step, it raises rather than quietly running dense and calling it a speedup.

The errors are the troubleshooting

This node has no settings, so its value shows up when things break:

• "requires a model patched by Flux2JiTApply" - you fed SamplerCustom a model that didn't come out of the Apply node. The two nodes are a matched set; reconnect the patched MODEL.
• "does not support masked/inpaint denoising" - inpainting is out of scope for this pack. Don't pass a denoise mask.
• "stage logic is active, but the JiT diffusion-model wrapper did not execute" - the patched model got recloned, replaced, or bypassed somewhere between Apply and SamplerCustom, so no sparse inference is happening and there's no speedup to show. Check that SamplerCustom receives the Apply node's MODEL output directly and that no later node swapped it.

Using it

Load a Flux.2 model, patch it with Flux2 JiT Apply, generate sigmas with Flux2 JiTScheduler (or core Flux2Scheduler), then:

Flux2 JiT Apply (MODEL) ──> SamplerCustom (model)
Flux2 JiTScheduler (SIGMAS) ──> SamplerCustom (sigmas)
Flux2 JiT Sampler (SAMPLER) ──> SamplerCustom (sampler)

Start from an empty Flux.2 latent and decode normally. Keep it to one-eval-per-step samplers - the README is explicit that multi-eval paths like some Heun/DPM-family samplers are not supported. To confirm the speedup is real, turn on verbose in the Apply node once; the sampler prints a wrapper summary with sparse vs dense call counts.

Installing it

Same story as the rest of the pack - ComfyUI Manager search for "JiT Flux2", or:

cd ComfyUI/custom_nodes
git clone https://github.com/xmarre/ComfyUI-JiT-Flux2.git

then restart ComfyUI. No extra Python dependencies and no model files to fetch; it runs on the Flux.2 checkpoint you already have. Flux.2 is a heavy model - Dev alone is ~35GB fp8 plus its encoder - so "faster for free, no retraining" is a real offer here. Just don't expect it to fix OOMs (it cuts compute, not VRAM), and don't expect bit-identical output to a dense run. JiT trades a hair of exactness for wall-clock time, and that's the whole deal.