Read Attention Maps (Experimental)

A recording of attention is a tensor. A tensor is not an answer. Read Attention Maps turns the stack Attention Map Capture recorded into an IMAGE batch - a heat video where bright means "this word is looking here" - so you can watch a prompt word walk across your frames instead of squinting at numbers.

It renders nothing through the VAE and it doesn't change your generation. It's a viewer.

How it works

The capture stored maps keyed by (step, block), each a stack of frames on the latent grid, one per recorded token. This node:

1. picks your token - sum across every recorded token, or one token index;
2. picks your blocks - mean over all recorded blocks, or one block index;
3. picks your steps - all, mean, or one step index;
4. normalises (see below), upscales from the latent grid to the pixel grid with nearest-neighbour at the model's own patch × latent downscale factor, and colours it with an inferno ramp (or plain gray).

The frame count is what matters for playback. match_video_frames (on by default) repeats each latent frame four times after the first, because WAN's 3D VAE compresses time 4:1 - one latent frame is four video frames. With it on, your attention video lines up frame-for-frame against the decoded clip, which is the whole point: drop both into a video combine and compare.

Inputs, in the order you'll fight with them

• attention_maps - from Attention Map Capture.
• latent - the sampler's LATENT output, so this node runs after sampling. Same sequencing trick as the activation probe reader.
• token - sum (default) or a token index. If the capture had clip and prompt, the report will list which index maps to which word. sum is the "all words together" view and is often the most legible.
• block - mean (default) or a block index. If you want to compare mid-blocks against late blocks, read the same capture twice with different block values.
• step - all renders every recorded step side by side in one frame, which is a beautiful way to watch attention resolve from blob to object over the denoise. mean averages them. Or name a step index for a single panel.
• normalize - per_video (default) uses one scale for the whole clip, so frame-to-frame brightness changes are real. per_frame rescales every frame independently: better contrast per frame, but you can no longer read relative intensity across time.
• colormap - inferno or gray.
• match_video_frames - as above.

Outputs: frames (IMAGE batch → Video Combine, SaveWEBP, Preview Image) and report, a JSON string listing the recorded steps, blocks, token labels and how many frames came out.

Practical loop

Capture on the model chain, sample, then read. If you're comparing an unbent run to a bent one, the cleanest setup is a single workflow with two capture + read pairs, matched on the same blocks and tokens - which is exactly what the shipped wan21_t2v_attention_capture workflow does: capture horse, then rotate the attention 20° in the early steps, then capture horse again downstream, and read both as video.

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

No extra weights - kornia and scikit-learn are the whole dependency list. Restart ComfyUI. Everything is logged under the [model-bending] tag.

Where people get burned

• "No attention maps were recorded." The capture's MODEL didn't reach the sampler, or this node's latent isn't wired to a sampler that used it. Both wires, every time.
• A step or block index that wasn't recorded. The node errors and lists what was recorded - genuinely useful, and a sign your capture spec and your read spec drifted apart.
• Expecting smooth footage. These maps are one cell per latent cell, upscaled nearest-neighbour. They're supposed to be blocky; interpolating them would be inventing information.
• Watching the sum view and concluding a word is diffuse. sum includes padding keys and every other token. Record a specific word if you want to know about that word.
• Assuming a clean map means a clean image. Attention tells you where the model looked, not what it drew. They usually correlate, and when they don't, that's the interesting part.