Gen2 Flux.2 Klein Fix (#12905 revert)

Did you upgrade ComfyUI to v0.17.0 or later and suddenly your Flux.2 [klein] masked-inpaint results look different - same workflow, same seed, different image? It wasn't your imagination, and it wasn't a bad checkpoint. A core refactor (commit 44f1246, PR #12905, "Support flux 2 klein kv cache model: Use the FluxKVCache node") changed how Klein's transformer runs for normal index reference-latent workflows like masked inpaint, even though the new KV-cache code paths it introduced are gated off for those exact workflows. This node is the targeted revert for that one regression.

Klein is the default local editor for a lot of people in 2026 - small, fast, does single- and multi-reference editing in one checkpoint. The kv-fp8 speed trick (the FluxKV node) is genuinely worth it for multi-reference edits. But if you never opted into KV caching, you still got its side effects. That's the part this node exists to undo.

How it works

It's about as surgical as a compatibility patch gets. The node takes one model input, clones it, and installs a unet_function_wrapper that swaps the flux transformer's forward_orig and _forward back to the last-good v0.16.4 implementations only during this model's own sampling, then restores them immediately after. The vendored functions are copied verbatim from comfy/ldm/flux/model.py at the commit right before 44f1246. No ComfyUI core files are edited, so the patch survives git pull - which is the whole reason it exists, because the alternative (hand-editing model.py) dies on every upgrade.

Wiring and the one trap

The input and output are both MODEL, and there's nothing else to configure. Wire it in series like this:

Load Diffusion Model → Gen2 Flux.2 Klein Fix → guider

That's the entire setup. The node does nothing until sampling runs, and it only changes behavior on the Klein path it's reverting.

The trap is the one thing the README warns about: don't combine this with FluxKVCache or the newer flux2 edit (index_timestep_zero) features. Those are exactly what this node reverts. If you're on a kv-fp8 Klein checkpoint and using the KV node for multi-reference speed, you do not want this fix on the same model - you'd be re-adding the pre-cache path and throwing away the speedup. This is for the plain index reference-latent / masked-inpaint case, nothing else.

Installation

Install the pack once via ComfyUI Manager (search "ComfyUI-gen2"), or manually:

cd ComfyUI/custom_nodes
git clone https://github.com/petmycat/ComfyUI-gen2
cd ComfyUI-gen2
pip install -r requirements.txt

Then restart ComfyUI. This section has no extra dependencies - the only thing requirements.txt pulls is scipy, which the tiling nodes use, and the pack loads sections independently anyway.

Common issues

• No change in output. Then you either aren't on v0.17.0+, or your workflow isn't hitting the affected path (plain text-to-image on Klein was never the problem - reference-latent / inpaint is). The node is a pass-through otherwise, so there's nothing to misconfigure.
• You're using KV caching. See above - if you want the KV speedup, leave this node out.
• The whole area was churny at launch. The KV-cache node itself shipped with OOM reports that ComfyUI had to hotfix (issue #12906, fix PR #12909). If you're seeing memory errors, that's the KV path, not this node - and another reason to run the reverts and the caches on separate models.

One node, one job, done right. If your Klein inpaints changed after an upgrade and you don't use KV caching, this is likely the cleanest fix you'll find - no core edits, no speed cost, no config.