Flux2 JiT Scheduler

The most exciting node in a speedup pack is never the scheduler. Nobody posts screenshots of a SIGMAS curve. But in ComfyUI-JiT-Flux2, Flux2JiTScheduler is the node that keeps the whole trick from drifting off course - because JiT's entire bet is that the sampling trajectory is a straight line you can afford to evaluate sparsely. Feed it wrong sigmas and you're accelerating garbage.

What it is: a drop-in replacement for ComfyUI's core Flux2Scheduler, built for the JiT custom-sampling path. It takes your step count and image dimensions, computes Flux.2's empirical sequence-length shift from those dims, and hands you a SIGMAS tensor for SamplerCustom. In the workflow the README prescribes, it sits between Flux2 JiT Apply (which patches the model) and Flux2 JiT Sampler (which rides inside SamplerCustom). JiT is training-free - this node isn't loading special weights, it's computing the noise schedule the sparse sampler needs.

How it works

Flux.2 is a rectified-flow model, and flow-matching models use a timestep shift to spread denoising effort between composition and fine detail. That shift isn't a fixed constant; ComfyUI derives it from how many latent tokens the image actually produces. That's what this node does: it quantizes your image-space dimensions down to the latent token grid (a 1024×1024 image is a 64×64 grid), computes the sequence length, derives the empirical mu for your step count, and applies the generalized time-SNR shift to the base linear schedule. In flux2 mode that's effectively the Flux2Scheduler path you already know, reimplemented for the JiT graph.

Inputs that matter

Seven inputs, and you'll actually touch three:

• steps - default 18. Match this to the preset on Flux2 JiT Apply: 18 for default_4x, 11 for default_7x. This is the single most common mismatch.
• width / height - final image-space dimensions, default 1024×1024. Do not pass latent-token sizes. Non-16-aligned values are quantized internally to the token grid before the shift is computed, so you don't need to round anything yourself.
• schedule - flux2 (default) or jit_beta. The beta mode applies a JiT-paper-style beta warp before the Flux.2 shift. The README is blunt: treat it as experimental until someone visually validates it on Flux.2. It's the FLUX.1-dev paper schedule bolted onto a different model. Leave it on flux2 unless you're specifically chasing the paper.
• beta_alpha (1.4), beta_beta (0.42), beta_resolution (8192) - only read in jit_beta mode, defaults from the paper's reported beta schedule. You will not touch these.

Output is a single SIGMAS that feeds the sigmas input of SamplerCustom.

Installing it

The pack installs like any custom node. Either search ComfyUI Manager for "JiT Flux2" / "ComfyUI-JiT-Flux2", or clone the README's canonical way and restart ComfyUI:

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

There's no requirements.txt - it leans on torch and einops, which ComfyUI already ships - and no model downloads. It works on whatever Flux.2 checkpoint you already run. The author (xmarre, "marres" on Reddit) has a track record of training-free acceleration ports like ComfyUI-Spectrum-WAN-Proper, which is worth knowing because it means the numbers in the README are benchmarked, not vibes.

Where people get burned

• Steps vs preset mismatch. The scheduler's steps and the Apply node's preset are separate numbers. Run default_7x (11 steps) with the scheduler left at 18 and you've paid for a longer run without the corresponding preset. The Apply node resets its expected steps from the preset; the scheduler just takes what you give it.
• Latent sizes. Some Flux workflows habitually pass token-grid dimensions. This node wants the real image size.
• jit_beta looks tempting because "paper-style" sounds superior, but the author flags it as unvalidated on Flux.2. Start with flux2. If you do experiment, A/B against a plain run at the same seed before trusting it.

One more honest caveat: JiT changes the actual ODE trajectory, so a JiT run won't be pixel-identical to a plain Flux.2 run at the same seed - that's the acceleration tax, and it's tiny. If you need bit-identical output, you don't want speedups. If you're tired of waiting on a 32B Flux.2 Dev, this is the honest way to go faster without downloading anything new.