Krea2 Prompt Reinjection

Deep transformers have a bad habit: by the time image tokens reach block 20 of 28, the text they're attending to has drifted a long way from what your prompt actually said. The detail you cared about - a specific colour, a specific object count, the word in quotes waiting to be rendered as text - loses the argument to whatever the image has already committed to. People call it prompt forgetting, and on Krea 2 with its Qwen3-VL encoder and long natural-language prompts it's the reason a two-sentence prompt sometimes lands as a one-sentence prompt.

Prompt Reinjection is the published fix for that: keep a copy of the shallow text features and add a little of them back before the deeper blocks. This node is the Krea 2 adaptation, and its defaults come from the official recommendation for a new MMDiT architecture rather than being guessed - origin_layer=1, targets 2..27, weight=0.025, no anchoring.

How it works

The node hooks Krea 2's block loop. At origin_layer it snapshots the positive text tokens (detached, so no gradient or memory weirdness - this is inference only). From target_start through target_end, before every block runs, it nudges the text tokens with target + weight * origin. Same treatment for the negative text stream if NAG is in the chain, using the negative snapshot from the same origin layer.

anchoring switches that plain add for the paper's optional token-wise anchoring: both tensors are mean/centred and standardised per token, mixed, layer-normalised and re-scaled back. It's off by default, which is what the official recommendation asks for on a fresh architecture, and it's the more invasive of the two - one more thing that can move your image, one more reason to leave it alone until you've seen a reason not to.

Because the snapshot is taken inside the same loop, origin_layer has to come before target_start, and the node validates that up front: origin_layer < target_start <= target_end < 28, else a ValueError. Layers are 0-27 because Krea 2 has 28 single-stream blocks.

Output is model (MODEL), same as the rest of the pack, so it drops into the chain between your checkpoint loader and your sampler. Order relative to GAG/NAG/TaylorSeer doesn't matter.

The inputs you'll actually touch

• weight (0.025, range -1 to 1) - the size of the nudge, and it really is this small. A twenty-five-thousandths add per block across 26 blocks is the point: death by a thousand papercuts, not a rewrite. Negative values push away from the shallow features instead.
• target_end (27) - how deep the injection runs. Drop it if you want the effect confined to the middle of the stack.
• anchoring (false) - as above; try it only if the plain version does nothing and you want to find out.
• origin_layer (1) - where "shallow" is defined. One is already near the input; there's not much reason to move it.

That's all five inputs, and four of them are integers with an obvious meaning. This is the least scary node in the pack.

Install

cd ComfyUI/custom_nodes
git clone https://github.com/iljung1106/Krea2-Training-Free-Improvements
# restart ComfyUI

Nothing to pip install - the pyproject declares zero dependencies, which is unusual for a guidance node and worth appreciating. ComfyUI Manager search "Krea2 Training Free Improvements" also works. Restart afterwards or the class won't be registered.

Where the wheels come off

The weight widget has a step of 0.001, so the field accepts 0.025 fine but the arrows jump a thousandth at a time. Hand-type it. Large values do not work better - you're overwriting the text features the model spent 26 blocks refining.

The node needs Krea 2's SingleStreamDiT and errors out on anything else, so don't expect to bolt it onto Flux and see prompt adherence improve. It also isn't a scaling knob for prompt strength; if your prompt is being ignored at step one, this is the wrong tool, because there's nothing to reinject yet.

And set your expectations accordingly: the author's own measurement puts the whole four-node stack (GAG + NAG + Prompt Reinjection) at 105.7 s on a 3060 Ti 8 GB versus 67.5 s for the acceleration path. Reinjection's contribution inside that is small - it's a two-line residual per block, not a second pass. Test it at a fixed seed, because the difference between "subtle improvement" and "I imagined it" is exactly the size of effect you're looking for.