Concept Attention

What it actually is

The name is doing real work. Concept Attention runs your loaded diffusion model once and tries to answer a surprisingly concrete question: which pixels in this image correspond to "cat", or "white", or "cane"? Not by loading a segmentation model, not by training anything - straight from the attention mechanism of the transformer that made the image.

That's the thesis of the ConceptAttention paper this pack ports: multi-modal diffusion transformers (DiTs) don't just generate images, they encode concepts in their attention output space, and a simple linear projection can pull those concepts back out. Flux, SD3/3.5, and CogVideoX all qualify. If you've ever wanted to know what the model "thinks" a word means inside a given image, this is the node that shows you.

Why you'd reach for it

Zero-shot localization without extra machinery. Instead of wiring up SAM or a dedicated segmenter (the KB's background-removal essay is the full story on how deep that hole goes), you hand the model a concept list and it tells you where those concepts live in the picture. That's useful for interpretability rabbit holes, for debugging why a prompt produced something odd, and as a stepping stone to masks for inpainting or background work.

How it works (and where the paper meets reality)

The paper's core claim is that a linear projection of the attention output gives sharper maps than raw attention - no gradient tricks, no training. This ComfyUI port does a simplified version: each concept goes through your CLIP encoder, the model runs forward, and the output becomes a per-concept map.

Here's the honest part, because it should set your expectations. This is a skeleton port, not a faithful reproduction. Reading the shipped source, it takes real shortcuts - random projection weights in places, and the "concept maps" it builds are near-flat tensors with a bit of sinusoidal texture rather than the crisp attention maps you'd see in the paper's figures. Load this expecting paper-quality saliency and you'll be disappointed. Load it wanting to poke around the idea inside ComfyUI, or as a starting point for your own implementation, and it's fine. One commit, zero community footprint - research-grade, not production-grade.

The inputs that matter

The node wants a full setup: model (MODEL), clip (CLIP), image (IMAGE), then two text fields. The one you'll actually edit is concept_list - a comma-separated list like woman, cat, white, lines, cane. That becomes your per-concept maps. prompt is in the schema, but read the code and you'll find it's not actually passed anywhere in the shipped version; concept_list does the work. (The README also advertises num_inference_steps and seed inputs - those don't exist in the shipped node either. Trust the schema, not the README.)

Outputs

Two of them:

• concept_maps (CONCEPT_MAPS) - the bundle of per-concept maps. This is the one that matters: it feeds the pack's other three nodes.
• visualized_image (IMAGE) - a visualization of the maps, ready for a SaveImage.

How to install

Via ComfyUI Manager, search "Concept-Diffusion" (pack title ComfyUI-Concept-Diffusion). Or manually:

cd ComfyUI/custom_nodes
git clone https://github.com/Junst/ComfyUI-Concept-Diffusion
cd ComfyUI-Concept-Diffusion
pip install -r requirements.txt

Then restart ComfyUI. No model downloads - it reuses the checkpoint and CLIP already in your graph. The requirements (torch, transformers, diffusers, einops, scipy, matplotlib) are all things a ComfyUI box almost certainly has; version conflicts are the main risk, same as any custom node (see the ecosystem essay on dependency pain).

Supported models and gotchas

Flux dev/schnell, SD3/3.5, CogVideoX, other DiT-based models. The code assumes a Flux-style apply_model(model, timestep, context, y) call, so old UNet checkpoints won't fit. And concept names are matched exactly, case-sensitively, downstream - "Cat" won't hit "cat".

If the node comes back with empty maps or just echoes your input image, that's its exception fallback - check the ComfyUI console for the real error, usually a shape or apply_model mismatch. It's a young pack; the console is your friend.