LTX Attn — Grid Viz

Metrics Heatmap gives you one colored cell per head. Grid Viz goes further - it renders an actual attention map per head, in a block×head grid, so instead of a single number per cell you get the full spatial pattern. It's the "stop squinting at one block at a time, look at the whole organism" node. If you're doing attention research on LTX-2.3's 1,536 heads, this is the overview you'll want pinned on a second monitor.

The layout mirrors the rest of the pack: X = blocks (0–47), Y = heads (0–31), with a separator line every 8th row and column so you can actually count cells. Each cell renders the key map, the query map, or their difference, depending on the view you pick.

How it works

It reads a capture store by handle and, for each targeted block and head, resolves the stored map at the chosen step, collapses it (same math as Key Map/Query Map - mean over queries or keys), reshapes to the frame layout, and renders it as a tile. What the node calls frame_mode controls how frames are handled:

• avg - average over all frames, one grid total.
• all - all frames stacked vertically inside each cell, one grid total.
• sequence - one grid per frame, output as an IMAGE batch.
• 0 or 3,7 - a specific frame index, or a list of them.

normalize decides the color scale: global (one scale across the whole grid - good for comparing relative strength), per_cell (each head normalized to its own max - good for seeing structure inside weak heads), per_block, or per_head. Default is per_cell, which is usually what you want on a first pass.

The inputs that matter

• store_handle - which capture to draw from (blank = current active store).
• view - key_map, query_map, or diff. diff is the block×head grid of key-vs-query difference, handy for spotting heads whose gaze and emission don't align.
• attn_type - sa or ca. Note the geometry implications: key maps only have spatial meaning for self-attention.
• target_blocks / target_heads - filters, default all. If RAM was tight on the capture you'll want to match what you actually recorded.
• step_idx - which denoising step; -1 = last captured.
• upsample, cell_padding, draw_labels, colormap - rendering knobs; upsample scales each low-res map up so cells are readable.

Outputs are grid_image (IMAGE) and size_info (STRING) - the size string tells you the resolved block/head count so you know the grid you're looking at is complete.

Install

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

Restart ComfyUI. No extra dependencies, no model files - same as the rest of this pack. It's work-in-progress, so store formats can change between versions.

Common issues

Grid Viz needs full maps, which means the capture used store_mode=full_fp16 or hybrid for the blocks you're viewing - a reduced capture has no dense map per head, so cells come back empty or the node skips them. Blocks captured via full_targets (sparse per-head) also break multi-head rendering here. Watch the two-run workflow too: the store is only populated after the capture generation finishes, so this runs in a later queue pass with the handle typed in. And at high cell_size/upsample with all 48 blocks, the output image gets genuinely large - if it chokes your preview, filter target_blocks down.