https://github.com/wooyeolbaek/attention-map-diffusers
🚀 Cross attention map tools for huggingface/diffusers
https://github.com/wooyeolbaek/attention-map-diffusers
attention-map cross-attention cross-attention-diffusers cross-attention-map diffusers huggingface stable-diffusion text-to-image visualization
Last synced: 6 months ago
JSON representation
🚀 Cross attention map tools for huggingface/diffusers
- Host: GitHub
- URL: https://github.com/wooyeolbaek/attention-map-diffusers
- Owner: wooyeolbaek
- License: mit
- Created: 2023-12-02T14:20:51.000Z (almost 2 years ago)
- Default Branch: main
- Last Pushed: 2025-01-18T18:24:38.000Z (9 months ago)
- Last Synced: 2025-04-08T12:09:04.360Z (6 months ago)
- Topics: attention-map, cross-attention, cross-attention-diffusers, cross-attention-map, diffusers, huggingface, stable-diffusion, text-to-image, visualization
- Language: Python
- Homepage: https://huggingface.co/spaces/We-Want-GPU/diffusers-cross-attention-map-SDXL-t2i
- Size: 7.89 MB
- Stars: 247
- Watchers: 2
- Forks: 19
- Open Issues: 8
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# Cross Attention Map Visualization
[](https://huggingface.co/spaces/We-Want-GPU/diffusers-cross-attention-map-SDXL-t2i)
Thanks to HuggingFace [Diffusers](https://github.com/huggingface/diffusers) team for the GPU sponsorship!
This repository is for extracting and visualizing cross attention maps, based on the latest [Diffusers](https://github.com/huggingface/diffusers) code (`v0.32.0`).
For errors reports or feature requests, feel free to raise an issue.
## Update Log.
[2024-12-22] It is now compatible with _"Stable Diffusion 3.5"_, _"Flux-dev"_ and _"Flux-schnell"_! (“Sana" will be the focus of the next update.)
[2024-12-17] Refactor and add setup.py
[2024-11-12] _"Stable Diffusion 3"_ is compatible and supports _batch operations_! (Flux and "Stable Diffusion 3.5" is not compatible yet.)
[2024-07-04] Added features for _saving attention maps based on timesteps and layers_.
## Compatible models.
Compatible with various models listed below.
- [black-forest-labs/FLUX.1-schnell](https://huggingface.co/black-forest-labs/FLUX.1-schnell)
- [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev)
- [stabilityai/stable-diffusion-3.5-medium](https://huggingface.co/stabilityai/stable-diffusion-3.5-medium)
- [stabilityai/stable-diffusion-3-medium-diffusers](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers)
- [stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
- [stable-diffusion-2-1-base](https://huggingface.co/stabilityai/stable-diffusion-2-1-base)
- ...
## Example.
cap-
-y-
-bara
hello
world
## demo
```bash
git clone https://github.com/wooyeolBaek/attention-map-diffusers.git
cd attention-map-diffusers
pip install -e .
```
or
```bash
pip install attention_map_diffusers
```
### Flux-dev
```python
import torch
from diffusers import FluxPipeline
from attention_map_diffusers import (
attn_maps,
init_pipeline,
save_attention_maps
)
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
torch_dtype=torch.bfloat16
)
# pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
pipe.to('cuda')
##### 1. Replace modules and Register hook #####
pipe = init_pipeline(pipe)
################################################
# recommend not using batch operations for sd3, as cpu memory could be exceeded.
prompts = [
# "A photo of a puppy wearing a hat.",
"A capybara holding a sign that reads Hello World.",
]
images = pipe(
prompts,
num_inference_steps=15,
guidance_scale=4.5,
).images
for batch, image in enumerate(images):
image.save(f'{batch}-flux-dev.png')
##### 2. Process and Save attention map #####
save_attention_maps(attn_maps, pipe.tokenizer, prompts, base_dir='attn_maps-flux-dev', unconditional=False)
#############################################
```
### Flux-schnell
```python
import torch
from diffusers import FluxPipeline
from attention_map_diffusers import (
attn_maps,
init_pipeline,
save_attention_maps
)
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-schnell",
torch_dtype=torch.bfloat16
)
# pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
pipe.to('cuda')
##### 1. Replace modules and Register hook #####
pipe = init_pipeline(pipe)
################################################
# recommend not using batch operations for sd3, as cpu memory could be exceeded.
prompts = [
# "A photo of a puppy wearing a hat.",
"A capybara holding a sign that reads Hello World.",
]
images = pipe(
prompts,
num_inference_steps=15,
guidance_scale=4.5,
).images
for batch, image in enumerate(images):
image.save(f'{batch}-flux-schnell.png')
##### 2. Process and Save attention map #####
save_attention_maps(attn_maps, pipe.tokenizer, prompts, base_dir='attn_maps-flux-schnell', unconditional=False)
#############################################
```
### Stable Diffusion 3.5
```python
import torch
from diffusers import StableDiffusion3Pipeline
from attention_map_diffusers import (
attn_maps,
init_pipeline,
save_attention_maps
)
pipe = StableDiffusion3Pipeline.from_pretrained(
"stabilityai/stable-diffusion-3.5-medium",
torch_dtype=torch.bfloat16
)
pipe = pipe.to("cuda")
##### 1. Replace modules and Register hook #####
pipe = init_pipeline(pipe)
################################################
# recommend not using batch operations for sd3, as cpu memory could be exceeded.
prompts = [
# "A photo of a puppy wearing a hat.",
"A capybara holding a sign that reads Hello World.",
]
images = pipe(
prompts,
num_inference_steps=15,
guidance_scale=4.5,
).images
for batch, image in enumerate(images):
image.save(f'{batch}-sd3-5.png')
##### 2. Process and Save attention map #####
save_attention_maps(attn_maps, pipe.tokenizer, prompts, base_dir='attn_maps-sd3-5', unconditional=True)
#############################################
```
### Stable Diffusion 3.0
```python
import torch
from diffusers import StableDiffusion3Pipeline
from attention_map_diffusers import (
attn_maps,
init_pipeline,
save_attention_maps
)
pipe = StableDiffusion3Pipeline.from_pretrained(
"stabilityai/stable-diffusion-3-medium-diffusers",
torch_dtype=torch.bfloat16
)
pipe = pipe.to("cuda")
##### 1. Replace modules and Register hook #####
pipe = init_pipeline(pipe)
################################################
# recommend not using batch operations for sd3, as cpu memory could be exceeded.
prompts = [
# "A photo of a puppy wearing a hat.",
"A capybara holding a sign that reads Hello World.",
]
images = pipe(
prompts,
num_inference_steps=15,
guidance_scale=4.5,
).images
for batch, image in enumerate(images):
image.save(f'{batch}-sd3.png')
##### 2. Process and Save attention map #####
save_attention_maps(attn_maps, pipe.tokenizer, prompts, base_dir='attn_maps', unconditional=True)
#############################################
```
### Stable Diffusion XL
```python
import torch
from diffusers import DiffusionPipeline
from attention_map_diffusers import (
attn_maps,
init_pipeline,
save_attention_maps
)
pipe = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
)
pipe = pipe.to("cuda")
##### 1. Replace modules and Register hook #####
pipe = init_pipeline(pipe)
################################################
prompts = [
"A photo of a puppy wearing a hat.",
"A capybara holding a sign that reads Hello World.",
]
images = pipe(
prompts,
num_inference_steps=15,
).images
for batch, image in enumerate(images):
image.save(f'{batch}-sdxl.png')
##### 2. Process and Save attention map #####
save_attention_maps(attn_maps, pipe.tokenizer, prompts, base_dir='attn_maps', unconditional=True)
#############################################
```
### Stable Diffusion 2.1
```python
import torch
from diffusers import DiffusionPipeline
from attention_map_diffusers import (
attn_maps,
init_pipeline,
save_attention_maps
)
pipe = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-1",
torch_dtype=torch.float16,
)
pipe = pipe.to("cuda")
##### 1. Replace modules and Register hook #####
pipe = init_pipeline(pipe)
################################################
prompts = [
"A photo of a puppy wearing a hat.",
"A capybara holding a sign that reads Hello World.",
]
images = pipe(
prompts,
num_inference_steps=15,
).images
for batch, image in enumerate(images):
image.save(f'{batch}-sd2-1.png')
##### 2. Process and Save attention map #####
save_attention_maps(attn_maps, pipe.tokenizer, prompts, base_dir='attn_maps', unconditional=True)
#############################################
```