NS Cross-Attention Map Recycler
Reuse your early attention maps to keep the composition honest
- model
- MODEL
Here's the thing nobody tells you about diffusion: the first few steps decide where everything goes. The early cross-attention maps - the ones that say "this token lives in the top-left, this one in the bottom-right" - set the composition, and later steps mostly refine what's already laid out. NS Cross-Attention Map Recycler exploits exactly that. It collects those early attention maps, then re-injects them into later steps so the model keeps honoring its own original layout instead of drifting.
Why is that useful? Two concrete reasons, depending on your goal. First, compositional stability: if your generations keep having elements wander or fuse mid-run, forcing later steps to reuse the early layout pins the scene down. Second, consistency across generations: the node's reset_cache toggle literally has "more consistency" in its tooltip - turn caching off and it keeps the maps from the previous generation, so repeated runs land closer together. That's the same family of trick as attention-based control (think of what ControlNet does for structure, just via attention instead of a side network), but self-contained in one node with zero extra models.
How it works
This is a MODEL patch, so it slots between your Load Checkpoint and your KSampler and wraps the UNet's cross-attention. The lifecycle has three phases:
- Collection (steps up to
collection_end, default 0.3): the node computes attention normally but stashes the conditional side of each cross-attention map, averaged and blended with an exponential moving average (collection_momentum, default 0.7 - higher means earlier maps carry more weight). - Neutral (between collection and injection): attention runs normally, fast path, no overhead.
- Injection (from
injection_start, default 0.3): each step's attention map gets blended with the saved one, atinjection_strength(default 0.3), with an optionaldecaythat fades the injection as the run progresses so the end of sampling stays free to polish.
target_layers lets you scope all of this: all, mid_only, up_only, or down_and_mid (the up-sampling blocks are where high-frequency detail gets re-asserted, so up_only is the surgical option if full-patch feels heavy).
The inputs that matter
You'll realistically touch five of them: collection_end (0.05–0.5, how long to gather), injection_start (0.1–0.8, when to start reusing), injection_strength (0–1, how hard to blend), decay (0–1), and reset_cache (enable = clear per generation for diversity; disable = reuse across generations for consistency). The rest are the advanced dials - collection_momentum, target_layers - that you tune after you see what the defaults do.
Output is a MODEL, which you chain from the checkpoint into your sampler. No latent or conditioning changes.
What to expect
The effect is subtle at low injection_strength - you're nudging, not teleporting. Push strength up and compositions get stickier, but too high and you'll feel the image "locked in" early, with less freedom in the final steps. Start at 0.3, watch what breaks, adjust. There's also a compute cost during collection and injection phases (the node falls back to a fast attention path when nothing's being collected or injected, so the neutral phase is cheap).
Installing it
It's one node in ComfyUI-NS-Util:
cd ComfyUI/custom_nodes
git clone https://github.com/NakamuraShippo/ComfyUI-NS-Util
or ComfyUI Manager → Install via Git URL with the same repo URL, then restart. No models to download, dependencies are the pack's usual light set.
The honest take
Attention-map recycling is a genuinely interesting idea that most people won't reach for, because the fix it offers - "my composition drifts" - is usually handled by better prompting or a stronger CFG, not by patching attention. But when you've hit the wall where the composition is almost right and just wobbles, this is a far cheaper and more targeted tool than rebuilding the workflow. It's the kind of node you forget you have until you need it, and then it's exactly right.
Inputs (8)
| Name | Type | Default | Description |
|---|---|---|---|
| model | MODEL | — | |
| collection_end | FLOAT | 0.300.05–0.5 | — |
| injection_start | FLOAT | 0.300.1–0.8 | — |
| injection_strength | FLOAT | 0.300–1 | — |
| decay | FLOAT | 0.500–1 | — |
| target_layers | COMBO | 4 options: all, mid_only, up_only, down_and_mid | |
| collection_momentum | FLOAT | 0.700–1 | EMA momentum for collected attention maps. Higher = more weight on earlier maps |
| reset_cache | COMBO | enable | Enable: clear cached attention maps every generation (more diversity). Disable: reuse maps across generations (more consistency) |
Outputs (1)
| Name | Type | Description |
|---|---|---|
| MODEL | MODEL | — |