Krea2 Geometry-Aware Attention Guidance

Be honest about what you're installing here. This pack is a single-commit repo from August 2026, and searching Reddit for "geometry-aware attention" returns literally nothing - no threads, no ports, no before/afters. The fancier sibling nodes (NAG, Prompt Reinjection, TaylorSeer) are techniques with years of community use behind them. GAG is a paper port with one tested default. That's not a reason to skip it, but it is a reason to test it at a fixed seed rather than trusting it because the README has a table in it.

What it's for: Krea 2 is a single-stream transformer with a 12B parameter budget and a Qwen3-VL text encoder, and image patches reading the prompt is the single most direct lever you have on adherence. Softmax attention spreads that read out; this node sharpens it.

How it works

For image-query → positive-text retrieval only, the node computes the usual dense softmax distribution and a sparse one using Entmax with alpha=1.5. Entmax picks a support and zeroes the tail, so the sparse distribution is decisive where softmax is smeared. Then it takes the paper's parallel-only GAG correction: take the residual between sparse and dense, project it onto the sparse direction, cap its magnitude at eta, scale the result by guidance_scale, and finally blend back toward dense by strength.

The two implementation details I'd actually point at are in the source. First, image self-attention is left completely dense - GAG touches one partition of the attention matrix, not the whole thing. Second, the correction is multiplied by the fraction of softmax probability mass that the query actually assigned to text keys, computed chunk-wise so a full QxK matrix never gets materialized. In other words it only pushes where there was already text attention to push on.

One caveat the author states plainly and you should believe: the paper proves GAG for explicit cross-attention, and Krea 2 uses joint attention. So this preserves native text-attention mass and replaces the conditional distribution inside the positive-text partition - an adaptation, not the authors' official code.

The inputs that matter

model, then this:

• strength (default 0.6) - the blend between native dense attention and the guided one. If output looks overcooked, this is the first thing down. The README's advice is exactly that: defaults follow the paper, Krea 2 is a different architecture, so reduce strength before anything else and compare against the same seed.
• guidance_scale (5) and eta (15) - how hard you amplify along the corrected direction, and the cap that stops a single query from exploding.
• sigma_start (0.881) and sigma_end (0.678) - the node only fires inside this sigma band. Default behaviour is mid-trajectory: structure gets established first, then retrieval gets sharpened. Outside the band, the wrapper skips the whole code path and you get a plain forward pass, which is why "I tried it and nothing changed" often just means "my schedule wasn't in the band."

Output is model (MODEL). Chain it into your sampler - the pack's convention is MODEL -> GAG -> NAG -> Prompt Reinjection -> TaylorSeer Lite -> KSampler, though the four MODEL nodes are independent and any order works.

Install

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

Zero dependencies declared - no pip install, no model download beyond Krea 2 itself. ComfyUI Manager will find it by searching "Krea2 Training Free Improvements"; the pack is on the Comfy Registry under publisher iljung1106.

Where people get burned

Use CFG = 1.0. The wrapper explicitly checks the conditioning branch list and raises "Krea2 training-free guidance requires CFG 1.0 (positive branch only)" if it sees anything but the positive pass - GAG is not a substitute for CFG, it's a different axis. The node also requires Krea 2's SingleStreamDiT and will refuse a Flux or SDXL checkpoint with a type error rather than silently doing nothing.

Do not stack it with legacy Krea2 NAG or Krea2Edit node packs; the loader detects those wrappers and errors out with the list of conflicts, which is a much nicer failure than a silent double-patch.

Expect a real time cost, not a free upgrade. In the author's four-node test on a 3060 Ti, GAG + NAG + Prompt Reinjection came in at 105.7 s versus 67.5 s for the order-0 TaylorSeer setup - guidance is the part that makes images nicer and the part that makes them slower.