wooyeolbaek/attention-map-diffusers

attention map tools for huggingface/diffusers

What it is

attention‑map‑diffusers is a Python package that lets you capture and visualise the internal attention tensors of diffusion models that run on the Hugging Face Diffusers library. It works for both image‑generation pipelines (e.g., Stable Diffusion, FLUX) and video‑generation pipelines (e.g., Wan, CogVideoX). By recording the cross‑attention and self‑attention matrices during inference, the tool produces heat‑maps that show how tokens or image patches influence each other – for example, which part of a prompt drives a particular region of the generated picture, or how a video frame attends to earlier frames.

Why it matters

Diffusion models are black‑box generators; understanding where the model is looking helps researchers debug prompts, study model behaviour, and develop better conditioning techniques. This library makes that inspection straightforward without modifying the original model code.

Core features (as described in the README)

Feature Details
Relation‑aware capture You can specify which attention relations to record, e.g., text→image, image→image, video→text, video→video.
Support for many models Pre‑tested with dozens of checkpoints: FLUX.2‑Klein, Stable Diffusion 3/XL, SD‑2, Z‑Image‑Turbo, CogVideoX‑2B, Wan‑2.1‑T2V‑1.3B, HunyuanVideo‑1.5, etc.
Simple API Use a context manager AttentionCapture(pipe, relations=…, offload=…) around a normal Diffusers pipeline call.
Automatic visualisation After capture, compute().save(...) writes PNG/GIF overlays and a JSON metadata file.
Video handling For video models, attention maps keep the (T, H, W) patch grid; the tool can render per‑frame PNGs and animated GIFs.
Resource‑aware Default guard of 8 GiB prevents RAM/VRAM overflow; you can tune max_capture_bytes and choose which timesteps to record.
CLI demos & validation demo/run_attention_demo.py runs quick examples for any supported model; test suite (pytest) and audit scripts verify coverage and parity with original generation.
Installation Single pip install: pip install attention-map-diffusers==1.0.0.
Licensing & citation MIT‑licensed; DOI and BibTeX entry provided for academic use.

How to get started (quick‑start snippets from the README)

Image example (FLUX.2‑Klein)

import torch
from diffusers import Flux2KleinPipeline
from attention_map_diffusers import AttentionCapture, text_tokenizers

prompt = "A red fox beside a glowing blue lantern in a snowy forest."
pipe = Flux2KleinPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-4B", torch_dtype=torch.bfloat16
).to("cuda")

with AttentionCapture(
    pipe,
    relations=["text->text", "text->image", "image->text", "image->image"],
    offload="cuda",
) as capture:
    images = pipe(prompt=[prompt], num_inference_steps=4).images

capture.compute().save(
    "outputs/attention", tokenizer=text_tokenizers(pipe),
    prompts=[prompt], images=images,
)

Video example (Wan‑2.1‑T2V‑1.3B)

import torch
from diffusers import WanPipeline
from attention_map_diffusers import AttentionCapture, VisualizationConfig, text_tokenizers

prompt = "A cinematic tracking shot of a red fox running across snow in a pine forest."
pipe = WanPipeline.from_pretrained(
    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", torch_dtype=torch.bfloat16
).to("cuda")
pipe.scheduler.set_timesteps(50, device="cuda")
last_timestep = [float(pipe.scheduler.timesteps[-1])]

with AttentionCapture(
    pipe,
    relations=["video->text", "video->video"],
    timesteps=last_timestep,
    relation_query_indices={"video->video": "center"},
    offload="cpu",
    max_capture_bytes=16*1024**3,
) as capture:
    videos = pipe(
        prompt=[prompt], num_inference_steps=50,
        height=480, width=832, num_frames=81,
    ).frames

# move modules back to CPU to free GPU memory
for comp in pipe.components.values():
    if isinstance(comp, torch.nn.Module):
        comp.to("cpu")
torch.cuda.empty_cache()

capture.compute(compute_device="cuda").save(
    "outputs/attention", tokenizer=text_tokenizers(pipe),
    prompts=[prompt], videos=videos,
    visualization_config=VisualizationConfig(max_items=16, video_fps=16),
)

What you get on disk

The package writes a directory tree like:

attention/
  metadata.json                # capture settings & provenance
  raw/<relation>/*.pt          # optional raw tensors (if --save-raw used)
  visuals/<relation>/aggregate/
    maps/*.png                 # per‑frame or per‑image heat‑maps
    maps/*.gif                 # animated GIFs for video relations
    overlays/…                 # original media with overlayed attention

Metadata includes the model checkpoint, prompt, timesteps captured, and any resource limits that were applied.

Who might use it

  • Researchers probing how diffusion models attend to prompt tokens or image patches.
  • Prompt engineers who want to see why a particular phrase influences a region of the output.
  • Educators demonstrating the inner workings of cross‑attention in generative AI.
  • Developers building debugging tools or visual‑explainability extensions for Diffusers pipelines.

All information above is taken directly from the repository’s README; no additional features have been inferred.

Related