什么是DAAM？利用交叉注意力解释Stable Diffusion

已更新以支持Stable Diffusion XL (SDXL) 和 Diffusers 0.21.1！

我会定期更新此代码库。如有任何问题，请提交issue。

在我们的论文中，我们提出了扩散注意力归因图（DAAM），这是一种基于交叉注意力的方法，用于解释Stable Diffusion模型。 请查看我们的演示：https://huggingface.co/spaces/tetrisd/Diffusion-Attentive-Attribution-Maps。 您还可以参阅由GitHub Pages托管的文档，以及为v0.1.0版本更新的Colab笔记本。

快速入门

首先，请根据您的平台安装PyTorch。 然后，您可以使用pip install daam来安装DAAM；如果您需要可编辑的库版本，则可以执行git clone https://github.com/castorini/daam && pip install -e daam。 最后，使用huggingface-cli login登录以获取大量Stable Diffusion模型——您需要在HuggingFace.co上获取一个token。

运行网站演示

只需在终端中运行daam-demo，然后访问http://localhost:8080即可。 您将看到与HuggingFace Spaces上相同的演示页面。

将DAAM用作命令行工具

DAAM自带一个简单的生成脚本，方便用户快速试用。 尝试运行以下命令：

$ mkdir -p daam-test && cd daam-test
$ daam "一只狗在田野上奔跑。"
$ ls
a.heat_map.png    field.heat_map.png    generation.pt   output.png  seed.txt
dog.heat_map.png  running.heat_map.png  prompt.txt

此时，当前工作目录下将包含生成的图像output.png、每个单词对应的DAAM热力图，以及一些辅助数据。 您可以通过运行daam -h查看更多选项。 若要使用Stable Diffusion XL作为后端，请运行daam --model xl-base-1.0 "狗在跳跃"。

将DAAM用作库

您可以按如下方式导入并使用DAAM：

from daam import trace, set_seed
from diffusers import DiffusionPipeline
from matplotlib import pyplot as plt
import torch


model_id = 'stabilityai/stable-diffusion-xl-base-1.0'
device = 'cuda'

pipe = DiffusionPipeline.from_pretrained(model_id, use_auth_token=True, torch_dtype=torch.float16, use_safetensors=True, variant='fp16')
pipe = pipe.to(device)

prompt = '一只狗在田野上奔跑'
gen = set_seed(0)  # 为了可重复性

with torch.no_grad():
    with trace(pipe) as tc:
        out = pipe(prompt, num_inference_steps=50, generator=gen)
        heat_map = tc.compute_global_heat_map()
        heat_map = heat_map.compute_word_heat_map('dog')
        heat_map.plot_overlay(out.images[0])
        plt.show()

您也可以非常方便地序列化和反序列化DAAM热力图：

from daam import GenerationExperiment, trace

with trace(pipe) as tc:
    pipe('一只狗和一只猫')
    exp = tc.to_experiment('experiment-dir')
    exp.save()  # experiment-dir 现在包含了所有数据和热力图

exp = GenerationExperiment.load('experiment-dir')  # 加载实验

我们将继续完善文档。在此期间，您可以查看GenerationExperiment、GlobalHeatMap和DiffusionHeatMapHooker类，以及daam/run/*.py中的示例脚本。 您还可以从论文附带的链接下载COCO-Gen数据集：http://ralphtang.com/coco-gen.tar.gz。 如果浏览器无法直接打开该链接，请将其复制到新标签页中，或使用wget等命令行工具。

另请参阅

• DAAM-i2i，DAAM在图像到图像归因任务上的扩展。
• Furkan的视频，介绍如何轻松上手DAAM。
• 1littlecoder的视频，展示了旧版DAAM的代码演示及Colab笔记本。

引用

@inproceedings{tang2023daam,
    title = "What the {DAAM}: Interpreting Stable Diffusion Using Cross Attention",
    author = "Tang, Raphael  and
      Liu, Linqing  and
      Pandey, Akshat  and
      Jiang, Zhiying  and
      Yang, Gefei  and
      Kumar, Karun  and
      Stenetorp, Pontus  and
      Lin, Jimmy  and
      Ture, Ferhan",
    booktitle = "Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)",
    year = "2023",
    url = "https://aclanthology.org/2023.acl-long.310",
}

DAAM 快速上手指南

DAAM (Diffusion Attentive Attribution Maps) 是一个基于交叉注意力机制的工具，用于解释和分析 Stable Diffusion 模型的生成过程。它能生成热力图，直观展示提示词中的每个单词如何影响最终图像的特定区域。

环境准备

在开始之前，请确保满足以下系统要求和前置依赖：

• 操作系统：Linux, macOS 或 Windows
• Python：建议 Python 3.8 及以上版本
• GPU：推荐使用 NVIDIA GPU 以加速推理（需安装对应的 CUDA 驱动）
• 前置依赖：
  • PyTorch：必须先安装与你的硬件环境匹配的 PyTorch 版本。
  • Hugging Face 账号：需要注册账号并获取 Access Token 以下载 Stable Diffusion 模型。

国内加速建议：

• 安装 PyTorch 时，建议访问 PyTorch 官网 选择对应命令，或使用清华/中科大镜像源。
• 下载模型时，若网络受限，可配置 HF_ENDPOINT 环境变量使用 Hugging Face 国内镜像：

  export HF_ENDPOINT=https://hf-mirror.com
  

安装步骤

1. 安装 PyTorch 请访问 PyTorch 官网 获取适合你平台的安装命令并执行。例如：

   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
   

2. 安装 DAAM 通过 pip 直接安装稳定版：

   pip install daam
   

   如果你需要修改源码或使用最新开发版，可以使用可编辑模式安装：

   git clone https://github.com/castorini/daam
   cd daam
   pip install -e .
   

3. 登录 Hugging Face 为了能够下载 Stable Diffusion 模型，需要在终端登录你的 Hugging Face 账号：

   huggingface-cli login
   

   按提示输入你在 HuggingFace.co 获取的 Access Token。

基本使用

DAAM 提供了命令行工具 (CLI) 和 Python 库两种使用方式。

方法一：使用命令行工具 (最简单)

这是快速体验 DAAM 功能的最快方式。运行以下命令将生成图像及对应的单词热力图。

# 创建测试目录并进入
mkdir -p daam-test && cd daam-test

# 运行生成命令
daam "A dog running across the field."

# 查看生成的文件
ls

输出说明： 当前目录下将包含：

• output.png: 生成的原始图像。
• *.heat_map.png: 提示词中每个单词（如 dog, running, field）对应的注意力热力图。
• prompt.txt, seed.txt: 生成参数记录。

使用 Stable Diffusion XL (SDXL)： 若想使用更新的 SDXL 模型作为后端，添加 --model 参数：

daam --model xl-base-1.0 "Dog jumping"

方法二：作为 Python 库使用

你可以在自己的代码中集成 DAAM 进行更灵活的控制和分析。

from daam import trace, set_seed
from diffusers import DiffusionPipeline
from matplotlib import pyplot as plt
import torch

# 配置模型和设备
model_id = 'stabilityai/stable-diffusion-xl-base-1.0'
device = 'cuda'

# 加载管道
pipe = DiffusionPipeline.from_pretrained(
    model_id, 
    use_auth_token=True, 
    torch_dtype=torch.float16, 
    use_safetensors=True, 
    variant='fp16'
)
pipe = pipe.to(device)

prompt = 'A dog runs across the field'
gen = set_seed(0)  # 设置随机种子以保证结果可复现

# 使用 trace 上下文管理器捕获注意力图
with torch.no_grad():
    with trace(pipe) as tc:
        out = pipe(prompt, num_inference_steps=50, generator=gen)
        
        # 计算全局热力图
        heat_map = tc.compute_global_heat_map()
        
        # 计算特定单词 ('dog') 的热力图
        heat_map = heat_map.compute_word_heat_map('dog')
        
        # 绘制叠加图
        heat_map.plot_overlay(out.images[0])
        plt.show()

保存与加载实验数据： DAAM 支持将生成的实验数据（包含所有热力图和元数据）序列化保存到磁盘，方便后续分析。

from daam import GenerationExperiment, trace

# 保存实验
with trace(pipe) as tc:
    pipe('A dog and a cat')
    exp = tc.to_experiment('experiment-dir')
    exp.save()  # 数据已保存至 experiment-dir 目录

# 加载实验
exp = GenerationExperiment.load('experiment-dir')