Nodes/ComfyUI Model Bending/Attention Map Capture (Experimental)
ComfyUI Node

Attention Map Capture (Experimental)

Watch where the word 'horse' actually looks

By abuzreq·Created 2 years ago·Updated 11 days ago· 26
Attention Map Capture (Experimental)
  • model
  • clip
  • MODEL
  • attention_maps
  • report
◄attentioncross_text►
◄blocks13-18►
◄tokensprompt►
◄steps*►
◄heads*►
◄prompt►
◄strictfalse►

Prompt adherence in video is a black box: you type "a horse galloping through a shallow river" and the model puts the horse somewhere, and you find out where by rendering 81 frames. Attention Map Capture tells you in advance - it records where each of your words is looking, frame by frame, and Read Attention Maps turns that into something you can watch.

It's the observational twin of Attention Map Bending. Same attention maps, no modification. Nothing about your video changes.

How it works

The node clones the model, and when attention is computed it records the cross-attention map for the tokens you named. It keeps the head mean (averaged across heads - that's the honest summary of "where does this token look") and only the first conditional sample, so a batch doesn't multiply the recording. It records at every step in your steps spec and every block in your blocks spec, and the maps come out on the latent grid, which the reader later scales up to pixels.

Under the hood it uses the same attention route as the bending node, so the recorded maps reflect everything upstream - put it after an Attention Map Bending and you're looking at the bent map, which is how the shipped capture workflow demonstrates a rotation actually doing something.

Inputs

  • model - goes on to the sampler as usual; the capture's model output is the one to sample with.
  • attention - cross_text (video-to-prompt) or cross_image (video-to-CLIP-image, for WAN 2.1 I2V). Those are the only two, because capture has nothing to say about self-attention.
  • blocks - default 13-18, the middle band of WAN 2.1 1.3B. On 14B, shift up to roughly 17–24.
  • tokens - default prompt, which records up to 32 prompt tokens at once. You can also name words (horse) if you wire clip and paste the prompt, or give indices like 0-3. Note the warnings: all is replaced with prompt on purpose because summing over every key always totals 1, which shows you nothing.
  • Optional: steps (default *, every step), heads (default *), clip, prompt, strict.

Outputs: MODEL, attention_maps (the recording, for Read Attention Maps), and report - a JSON string confirming model, blocks and token mode.

The wiring people get wrong

Two things have to be true at once:

  1. The capture's MODEL output feeds the sampler. Not the loaders' model - capture rewires the model chain, and if you plug the original in instead, nothing gets recorded.
  2. Read Attention Maps takes a latent input, and you connect the sampler's LATENT output to it. That's a sequencing trick: it forces the reader to evaluate after sampling instead of before.

Get either wrong and the reader throws "no attention maps were recorded" with a message explaining both wires.

What it costs

Only the tokens you asked for get kept, and only a head-mean, so a prompt capture over 32 tokens × 6 blocks × every step of a 1.3B run is fine. Push it - * blocks, every step, a 14B model - and you're storing a stack of feature maps. Record what you need, then read it.

Install

ComfyUI Manager → "ComfyUI Model Bending", or:

cd ComfyUI/custom_nodes
git clone https://github.com/abuzreq/ComfyUI-Model-Bending
pip install -r ComfyUI-Model-Bending/requirements.txt

kornia and scikit-learn, then restart ComfyUI. For the video nodes to actually do anything you need WAN family weights loaded normally - the capture doesn't ship models, and the example workflows in workflows/video/ link the ones they use.

Sanity checks

  • If you get maps that look like uniform noise, check the blocks. WAN's early blocks are broad and the late ones are near-uniform; the interesting structure lives in the middle band.
  • No clip and no prompt? Then tokens are recorded as indices and the report can't label them with words. It still works; you just have to guess which index was "horse".
  • Captured maps are a diagnostic, not a render. They'll be blocky by construction - one cell per latent cell, upscaled nearest-neighbour on purpose.
  • Marked experimental by the author, validated on tiny random-weight WAN models only. Reasonable to expect rough edges, not reasonable to file a bug about a pretty map.
Categorymodel_bending/video (experimental)

Inputs (9)

NameTypeDefaultDescription
modelMODEL—
attentionCOMBOcross_text2 options: cross_text, cross_image
blocksSTRING13-18Blocks to record, e.g. '13-18'
tokensSTRINGpromptTokens to record: prompt words 'horse' (needs clip + prompt), indices '0-3', or 'prompt' (each prompt token, at most 32)
stepsoptSTRING*Sampling steps to record
headsoptSTRING*—
clipoptCLIP—
promptoptSTRING—
strictoptBOOLEANfalse—

Outputs (3)

NameTypeDescription
MODELMODEL—
attention_mapsATTENTION_MAPS—
reportSTRING—