attention-map-diffusers records which encoder looked at which patch, and refuses to guess when it cannot

Thirteen checkpoints, and the relation list is not the same for all

The library's organising idea is the relation. Instead of one attention map per model, it names what is attending to what: text to text, text to image, image to text, and image to image, with the video pipeline adding text to video, video to text, and video to video. The example output is laid out as a two-by-two grid of exactly those four, with a token named in the corner of each cell.

The checkpoint table has thirteen entries and does not give all of them the same set. Several newer models are marked as supporting everything, which the file spells out as all four relations for an image pipeline and all four for a video pipeline, where the video set replaces image with video. Others, including the older Stable Diffusion family and one of the video models, are listed with a narrower pair. The reason a model appears in the table at all is that its denoiser exposes the attention modules the capture needs.

Requirements are Python 3.10 or newer and a recent version of the Diffusers library, with the floor stated as a specific release number rather than a range. Installation is one line, and the readme pins the exact version rather than leaving it floating:

pip install attention-map-diffusers==1.0.0

So picking a model key and picking which relations to request are two separate decisions, and the table exists to stop you requesting a relation a given architecture does not have.

The default keeps one timestep and a head mean, guarded at 8 GiB

The defaults here are conservative, and they are the right ones. Out of the box the final timestep and the head mean are retained, which means a first run gives you a small, cheap summary rather than the full trajectory. Asking for more is a flag away: one option captures every timestep, another saves maps for every layer instead of the aggregate.

The reason those flags are not on by default is stated plainly. There is a cumulative capture guard that defaults to 8 GiB, and its purpose is to prevent silent RAM and video memory exhaustion. That is a specific design choice: the tool would rather stop than hand you a half-written visualisation that looks complete. The guard is cumulative across the run, not per tensor, so a configuration that looks cheap per call can still add up.

The image path in the quick start is correspondingly small. A four-step generation on a 4-billion-parameter checkpoint, all four relations requested, and the capture is a context manager around the pipeline call, with the compute and save happening after the block exits. The video path adds an explicit scheduler setup and passes a list containing the last timestep value, so the timesteps argument is a list even when you want one.

There is also a shared runner that carries native per-model defaults and reproducible metadata, invoked once per model with a model key and a relations setting, so the per-model configuration lives in code rather than in your shell history:

python demo/run_attention_demo.py --model flux2-klein --relations all
python demo/run_attention_demo.py --model wan2.1-t2v-1.3b --relations all

Video self-attention would be a squared tensor, so it defaults to one patch

One default exists purely because the full computation is not practical, and the file explains the arithmetic. For image to image self-attention, every query patch is rendered by default, because the tensor is query patches against all patches and that is still a workable square. For video to video, the same relation would produce a tensor whose size is the query grid times itself in three dimensions, time, height, and width, and the project calls that impractical.

So video self-attention defaults to the center query patch. That is a single patch attending to every patch, which is a spatial map you can actually look at. The escape hatch is an explicit option accepting either a sampling specification or a request for everything, so the choice to spend the memory is yours to make rather than the library's to hide.

The video output has its own shape detail worth knowing. Maps preserve the patch grid as a three-dimensional time, height, and width structure, and each animated image also gets per-frame stills alongside it. The image and video examples in the readme are also presented differently on purpose: the spatial examples are overlays on the generated media, so the source image is not shown a second time, and spatial labels are off by default unless you ask for token or patch titles.

The readme justifies the whole layout with one observation about the video model used in the example, which exposes both cross-attention and spatial-temporal self-attention, so the table there shows every relation that model has.

Two text encoders can be one map or two, and the file says which

Token-axis concatenation is preserved, which sounds like a detail and is actually a correctness guarantee. When a pipeline has more than one text encoder, the capture keeps them apart along the token axis rather than merging them, and the two headline examples show the pattern. A model with a fused CLIP encoder and a T5 encoder saves both sets of maps separately, including all four component pairs of the text-to-text relation. A video model separates its two encoders after a mask-based token reordering step, and additionally records two of its internal refinement layers as a self-relation between the two halves of one encoder.

The honest counterpart is the other example given: for a model whose two encoders are fused along the feature axis rather than the token axis, the result stays one mathematically inseparable map. No flag splits it, and the file says so rather than shipping two images that are really one.

Classifier-free guidance adds a third wrinkle. When a model stores the guided and conditional passes as separate denoiser calls, each call is tagged, and the visualisations use the conditional call by default while the raw capture keeps both. So the map you see is the unguided one unless you ask otherwise, and that is a defensible default: the conditional pass is the one whose behaviour matches what the model was trained to produce.

Raw tensors are also optional and are never required for the visualisation, so a run that keeps only images is not a degraded run.

Capture fails with a layer report rather than returning half a result

The failure mode is the part of this library that most deserves copying. If any attention module in the denoiser is unsupported, capture fails and reports which layer and which processor was involved. The stated reason is to prevent silently incomplete results.

That is a real problem in this domain. Attention hooks are the kind of thing that attaches to most modules and quietly misses a few, and a visualization built from a partial set of layers looks exactly like one built from all of them. The only way a reader could tell the difference is if the tool refused to produce output at all. Here it does, and it tells you which layer to look at.

The same attitude shows up in the output layout. Every run writes a metadata file holding provenance and the capture settings, so a directory of maps carries the conditions that produced it. Raw per-relation tensors are written only when explicitly requested. Per-encoder maps and per-item frame directories sit in their own subdirectories, and generated-media overlays are kept separate from the attention maps themselves, so a map is never confused with a picture of the thing it was computed from.

The result is that a run is self-describing. If you find a map in a directory six months later, the metadata file next to it tells you the relation, the timestep policy, the model, and which encoders were involved. The layout makes the separation explicit:

attention/
  metadata.json                         # provenance + capture settings
  raw/<relation>/*.pt                   # only with --save-raw
  raw-components/<relation>/<encoder>/ # token-concatenated encoders
  visuals/<relation>/aggregate/batch-000/
    maps/*.png                          # image attention
    maps/*.gif                          # video attention
    maps/<encoder>/*                    # token-concatenated encoders
    maps/<item>/frames/*.png
    overlays/                           # generated-media overlays, separate

Fourteen demo scripts and four validators is how this project tests itself

The demo directory is the clearest statement of the project's structure, and it is organised one file per checkpoint. There are fourteen model-specific scripts, covering the 4-billion-parameter image model, a turbo image model, a pixel-style image model, both FLUX.1 variants, two Stable Diffusion 3 variants, two Sana and SDXL and SD2 entries, and three video models. Beside them sit a shared runner, a common module, and a coverage auditor.

One script per model is the tell. It means the per-model defaults live in code, not in a configuration file, which is what lets a single runner command produce native defaults and reproducible metadata for any of the thirteen keys. It also means adding a model is a code change with its own script, and that is a cost the project has accepted.

The four validators are where the quality claim is enforced. One checks generation parity, which is the important one: it verifies that attaching the capture does not change the image or video the pipeline produces. Without that check, every map in the repository could be describing a subtly different generation from the one you would get without the tool. Another validates the full matrix run, which is the cross product of models and relations. A third validates the option matrix, which is the cross product of the command-line flags. The fourth checks the attention outputs themselves, and the coverage auditor checks which attention modules were seen.

Unit tests are run as a plain pytest invocation, and the full matrix is also runnable in a dry-run mode that exercises the matrix without generating anything.

The readme on PyPI is a patched copy of the readme on GitHub

The packaging file contains the most human detail in the repository, and it is a detail about the readme itself. The long description sent to the package index is not the readme file: it is the readme read from disk and then rewritten. Relative asset paths are replaced with absolute raw-content URLs, because the package index has no repository to resolve a relative path against. The licence and citation links are rewritten to absolute repository URLs for the same reason.

So the two documents you read are not identical. On the repository the images are relative and the licence link points at a local file; on the package index the images are absolute and the licence link points at the web. Both render correctly, and neither is wrong, but a reader comparing them would find differences that are all mechanical.

The package metadata is otherwise plain and worth reading for the scope it declares. The name on the package index is hyphenated while the name you import is not. The version matches the newest release tag. Two authors are listed with academic addresses at two different universities. The description is one sentence naming the task and the library it works with. The keyword list covers attention maps, the library, diffusion models, interpretability, and both text-to-image and text-to-video.

The dependency file pins only three things: a floor on the Diffusers library, a floor on the transformer library, and a floor on the tensor library. Everything else is unpinned, including the tokenizer-related packages that a T5-based encoder needs, which is the one place where an unpinned dependency could plausibly change behaviour rather than just a bugfix level.

There is a citation file at the root, a manifest file, a space directory for the hosted demo, a test directory, and a licence. The hosted demo space is pointed at a specific model rather than at the library as a whole, so it is a worked example rather than a general playground.

Editorial conclusion

attention-map-diffusers is worth installing if you are debugging what a diffusion model actually looked at, because it separates the four relation types instead of collapsing them into one heatmap, and it tells you when two text encoders were fused into a single inseparable map rather than pretending otherwise. It is a poor fit if you want every timestep and every layer, because that is the configuration with a real memory cost and a guard that will stop you, and it is a poor fit if you need a video self-attention map for every query patch, which the file itself calls impractical. Before you read the output, check three things yourself: which relation you are looking at, since a text-to-text map and a text-to-image map answer different questions; which timestep you captured, because the default is the final one and the head mean, not the whole trajectory; and whether the map you are reading is a single fused map or a separated per-encoder one, since for some models the file says fusion is mathematically inseparable and no amount of configuration will split it.

Frequently asked questions

Official sources

1. License: MIT
2. Project website
3. README
4. Releases
5. wooyeolbaek/attention-map-diffusers on GitHub