FlowMatch Sampler (ai-toolkit exact)

If you've trained a Z-Image, Qwen-Image, or Flux LoRA in Ostris's ai-toolkit and then tried to sample it in ComfyUI, you've probably noticed the same thing everyone notices: it doesn't look like the previews your trainer gave you. Colors shift, the style weakens, you bump CFG and steps around and still can't get back to what you trained. That's not your LoRA. That's your sampler.

FlowMatch Sampler (ai-toolkit exact) is a single all-in-one node that swaps ComfyUI's default sampling for ai-toolkit's own flow-matching behavior - the model patch, the sigma schedule, the shift math - so what you trained with is what you get. The pack ships exactly one node, and it's this one.

Why the mismatch exists

This is a flow-matching problem, not a you problem. Modern models like Flux, Qwen-Image, and Z-Image are trained with the flow-matching objective: the model learns a near-straight path from noise to image, which is why they generate in 4–8 steps instead of 20–50. But "flow matching" isn't one thing - the exact sigma schedule and shift parameters differ per trainer. ComfyUI's built-in Flux sampler is close to ai-toolkit's schedule but not identical, and with distilled models the difference between "close" and "exact" is the difference between a clean Z-Image-Turbo output and a washed-out one.

How it works

First it patches the model's sampling config - flux and qwen get dynamic shift (shift scales with resolution), z-image gets the static shift=3.0 its training config expects. Then it builds the sigma schedule. Here's the clever bit: if it finds an ai-toolkit checkout sitting next to the custom node, it imports ai-toolkit's own CustomFlowMatchEulerDiscreteScheduler and uses that verbatim. If not, it falls back to local ports of the same formulas - so "exact" is exact when the toolkit is around, and very close when it isn't. Finally it runs the sample through a standard CFG guider with your chosen sampler name (default euler).

The preset defaults matter: steps=8, guidance_scale=1.0, width=768, height=1024, noise_seed=42 - these mirror the repo's own config.yaml sample block for Z-Image. Note that's distilled-model behavior. Sampling Z-Image Base with this means stepping up to 25–50 steps and CFG 3–5.

Inputs that actually matter

Most of the sixteen inputs you can just leave alone. The ones you'll touch:

• model_type - auto, flux, z-image, qwen. Defaults to z-image, and the README is blunt: switch it whenever you're sampling anything else. Forgetting this is the #1 way to get garbage, because the wrong shift is applied.
• steps and guidance_scale - start at the training config's values, not what feels right from SDXL muscle memory.
• sampler_name - euler or res_multistep. Pick whichever your training setup used; euler is the safe default.
• width / height - must match the resolution you're actually generating at, since dynamic shift is computed from them.
• force_aitk_timesteps - defaults on, and you want it on: it uses ai-toolkit's 1.0 → 1.0/steps timestep range before the shift math.
• shift / resolution_aware - leave at 0 / off until you know you need them; the presets already set sensible values.

The outputs: output is your denoised latent - wire that to a VAE Decode. denoised_output is a preview of the clean prediction (handy for img2img workflows), sigmas shows you the schedule it actually used, and patched_model is the model with the sampling patch, in case you want to chain another sampler off the same setup.

Installing it

ComfyUI Manager: search for ComfyUI-FlowMatch-Advanced. Or the manual route:

cd ComfyUI/custom_nodes
git clone https://github.com/supSugam/ComfyUI-FlowMatch-Advanced
cd ComfyUI-FlowMatch-Advanced
pip install -r requirements.txt

Then restart ComfyUI. Two things worth knowing before you do. The requirements.txt pins diffusers to a specific git commit so the sigma math can match ai-toolkit's build - that's an aggressive dependency, and if another custom node needs a different diffusers version you may have a conflict on your hands. And there are no model files to download; this works on the Flux / Qwen-Image / Z-Image checkpoints you already have. The node only makes sense for those architectures, so don't reach for it with SDXL or Illustrious.

Where people get burned

The pinned diffusers install is the one real tripwire, so if your environment is shared with other packs, keep an eye on what installing this overwrites. Beyond that: the z-image default will silently apply the wrong shift to Flux/Qwen models, and the node's "exactness" quietly degrades to "very close" if ai-toolkit's scheduler isn't discoverable on disk - still better than stock ComfyUI, just not bit-for-bit. It's a niche tool for a real annoyance: you trained in ai-toolkit, you want ai-toolkit out. For that job, nothing else ships this exact schedule.