Flux2 JiT Apply

Every other node in this pack is plumbing. Flux2JiTApply is the magic.

It's an unofficial ComfyUI port of Just-in-Time (JiT), a training-free spatial acceleration method for diffusion transformers (Sun, Li & Liu, arXiv 2603.10744). In plain terms: instead of running the whole latent token grid through the transformer at every denoising step, it evaluates a sparse subset of "anchor" tokens early on and progressively activates more as sampling proceeds - coarse-to-fine. That's why you'd reach for it: Flux.2 Dev is the local quality ceiling and also the reason people rejected it, a 32B rectified-flow transformer that needs ~35GB fp8 plus an ~18GB text encoder. JiT doesn't shrink weights or reduce VRAM; it cuts how much of the model you actually compute per step, which is wall-clock speedup with zero training and zero new files.

How it works

The node clones your model, attaches a config, and wraps the diffusion model with a sparse-token wrapper. During sampling the wrapper runs the transformer only on active tokens, lifts the sparse velocity back to the full grid (nearest-neighbor plus a controlled Gaussian blur that preserves exact anchor values), and lets newly activated tokens enter coherently via a deterministic micro-flow target rather than fresh noise. It's the paper's SAG-ODE + ITA + DMF stack. Stages run like 45% → 70% → 100% of the token grid - sparse early when the image is just shape, dense late when text and anatomy need every token.

Inputs that matter

The defaults are the point. You should change almost nothing here.

• model - Flux.2 only. It validates the model config and raises a ValueError otherwise, so no Flux.1, no other architectures. This is intentional, not a bug.
• preset - default_4x (18 steps, 45% → 70% → 100%) or default_7x (11 steps, 40% → 66% → 100%). The names are the expected speedup tier; the README is careful to call them "Flux.2-tuned conservative defaults," not the paper's more aggressive FLUX.1-dev numbers. Flux.2 needs higher early/mid token coverage than Flux.1 did, especially for text and small details.
• expected_total_steps - a hint, not a command. Runtime sigma count from the scheduler still controls execution; if they disagree, the sampler logs it and uses what it gets.
• stage_ratios / sparsity_ratios - comma-separated strings, only read in custom mode. Both lists must be the same length, at least two stages, ratios in (0,1], stage ratios strictly increasing and ending at exactly 1.0. The node errors rather than silently doing something weird, which is refreshing.
• use_checkerboard_init (on by default) - paper-style strided/checkerboard first stage with boundary coverage. Leave it on. The README's changelog is fun here: an earlier port made the non-checkerboard path "select the whole grid and randomly drop tokens," which visibly hurt quality. Checkerboard is default again for a reason.
• use_adaptive (on by default) - importance-guided token activation. Leave it on.
• microflow_relax_steps (3) - mostly a compatibility/control input; stage activation commits only after the DMF target is reached.
• blur_scale (0.4) - interpolation blur strength. Only touch it if interpolation looks wrong.
• verbose (off) - turn this on once, just to see your speedup. The sampler prints a summary of sparse vs dense vs fallback wrapper calls, which is the honest way to check JiT is actually firing.

Output is a single patched MODEL - a clone, so your original is untouched. Wire it into SamplerCustom's model input and don't let anything clone or bypass it afterward.

Installing it

Search ComfyUI Manager for "JiT Flux2" / "ComfyUI-JiT-Flux2", or clone and restart:

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

No requirements.txt, no model downloads - it uses torch and einops, which ComfyUI already has. The author, xmarre (marres on Reddit), does training-free acceleration ports with real benchmarks; this pack follows the same paper-first pattern.

Where people get burned

• Expecting VRAM relief. JiT cuts compute, not footprint. A Flux.2 checkpoint that OOMs will still OOM; one that fits will just render faster.
• Quality drift. JiT changes the actual ODE trajectory, so output differs slightly from a plain run. That's inherent, and the earlier port's two avoidable degradation paths (fresh transition noise, premature token commits) are fixed - but the trajectory difference remains. A/B at a fixed seed before you trust it for finals.
• No ControlNet, no inpainting. Those paths are explicitly out of scope. JiT is a straight text-to-image / image-generation path.