AttLine: Visualize Attention at a Line in Diffusers

Capture and visualize joint-attention maps from diffusion pipelines — per segment, per word, or per bounding box. Lightweight, no fork of diffusers needed.

Languages: English · 中文

Since diffusers>=0.35 reworked its attention dispatch interface, existing attention-visualization tools (which hook into the old attention-processor layer) no longer work on the new transformer stacks. Attline adapts to the new interface and wraps the whole thing behind a single line of code — attach(pipe, words=[...]) — so you can keep your normal pipeline usage completely unchanged and still get per-word heatmaps out of any supported pipeline, with minimal intrusion into your code.

Currently supported pipelines

Roadmap

• More pipeline integrations
• Image-editing model support
• bbox as a query selector

Install

1. Install diffusers from source or >=0.37.0:

pip install git+https://github.com/huggingface/diffusers.git

2. Clone and install AttLine:

git clone https://github.com/TeaWhiteBro/AttLine.git
cd AttLine
pip install -e .

Quickstart

The minimum to enable attention capture is a single line:

import torch
from diffusers import Flux2KleinPipeline
from attline import attach

pipe = Flux2KleinPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-9B", torch_dtype=torch.bfloat16
).to("cuda")

attach(pipe, words=["cat", "hello", "world"])

image = pipe(prompt="A cat holding a sign that says hello world").images[0]
image.save("out.png")
# Heatmaps saved in ./attn_out/

Call detach(pipe) to restore the original pipeline.

Usage — attach(pipe, ...)

Minimum required:

Optional:

Acknowledgement

This project is inspired by wooyeolbaek/attention-map-diffusers, which pioneered per-token joint-attention visualization for diffusion pipelines.

Citation

If you find this tool useful, please consider leaving a ⭐ on GitHub.

License

MIT. See LICENSE for details.