LTX Attn — Key Map

The LTX Attention Profiler pack's tagline for this node is "which tokens are being looked at?", and that's genuinely what it draws. Given a stored attention map for one block, it collapses the query dimension and reshapes the result into a spatial picture - so instead of a matrix of 14,000×14,000 numbers, you get a heatmap of the actual video frame showing where each head's attention lands. If you've ever wanted to see a model attend to a face or a moving hand rather than guess, this is the node.

It's self-attention only. The keys here are video tokens with real spatial geometry (1 token = 1 latent pixel, roughly a 32×32 patch in image space for LTX-2.3), so keys map cleanly onto frames. Cross-attention - video looking at text tokens - is Query Map's job instead.

How it works

From the stored full attention map W: [H, Sq, Sk] for a block, it computes W.mean(dim=1) - averaging over all query positions gives, for every key position, the total attention mass that lands there. That [Sk] vector is reshaped into [frames, latent_h, latent_w] and rendered as one colored tile per head, aligned to the frame layout. You get a grid where each cell is one head's "gaze map" over the video.

The inputs that matter

• block_idx - which of the 48 transformer blocks to visualize (0–47).
• head_indices - "all" or a comma list like "8,12,16" to narrow to the heads you care about.
• query_token_idx - -1 (default) averages over all query tokens, which is the "whole frame's gaze" view. A specific index isolates what one particular token looks at - that's the keypoint-tracking trick this pack is built around.
• aggregate_frames - how to collapse the frame dimension: mean, max, or first.
• num_frames, latent_height, latent_width - these must match the capture or the node raises. The sequence length has to equal frames × height × width exactly. Wire LTX Latent Dims into these (that's literally why that utility exists).
• step_idx - which captured denoising step to show; -1 is the last captured step.

Output is a single key_map IMAGE.

There's one prerequisite that bites people: this node needs the full attention map for the block, which means the capture ran with store_mode=full_fp16 (or hybrid with that block in full_blocks). reduced mode stores only the pre-collapsed key map, so there's nothing to resolve against. And if the block was captured via full_targets (per-head selection), you'll get an error - a sparse dict of heads isn't a dense multi-head map, and multi-head consumers like this one refuse it. Re-capture with the block in full_blocks instead.

Install

cd ComfyUI/custom_nodes
git clone https://github.com/g-raw/ComfyUI-LTX-Attention-Toolkit.git

Restart ComfyUI, and that's it - no extra dependencies beyond torch/numpy, no model files to download. One caveat: the pack is explicitly work-in-progress, so store formats can change and older .pt files may not load after an update. Pin the repo if you're building research workflows on it.

Common issues

The classic failure is the geometry mismatch - you captured at 16 frames and 22×40 latent, then left num_frames=1 and 32×32 here, and got Sq != T x Lh x Lw. Use LTX Latent Dims to extract the real numbers from your latent and feed them in. The second classic is trying to read maps from a reduced-mode capture: there's no dense map to show, so the node errors. And remember the two-run workflow - the store is only populated after the KSampler finishes, so run the capture generation, then queue this node with the handle typed in.