wooyeolbaek/attention-map-diffusers

attention map tools for huggingface/diffusers

是什么

attention‑map‑diffusers 是一个 Python 包，可让您捕获和可视化在 Hugging Face Diffusers 库上运行的扩散模型的内部注意力张量。它适用于图像生成管道（例如 Stable Diffusion、FLUX）和视频生成管道（例如 Wan、CogVideoX）。通过在推理过程中记录交叉注意力和自注意力矩阵，该工具生成热图，显示标记或图像块如何相互影响——例如，提示的哪一部分驱动了生成图像的特定区域，或视频帧如何关注之前的帧。

为什么重要

扩散模型是黑箱生成器；理解模型在何处关注有助于研究人员调试提示、研究模型行为，并开发更好的条件化技术。此库使这种检查变得简单，无需修改原始模型代码。

核心功能（如 README 所述）

功能 详情
关系感知捕获 您可以指定要记录的注意力关系，例如 text→image、image→image、video→text、video→video。
支持多种模型 已预先测试数十个检查点：FLUX.2‑Klein、Stable Diffusion 3/XL、SD‑2、Z‑Image‑Turbo、CogVideoX‑2B、Wan‑2.1‑T2V‑1.3B、HunyuanVideo‑1.5 等。
简单 API 在正常 Diffusers 管道调用周围使用上下文管理器 AttentionCapture(pipe, relations=…, offload=…)。
自动可视化 捕获后，compute().save(...) 会写入 PNG/GIF 叠加图和 JSON 元数据文件。
视频处理 对于视频模型，注意力图保持 (T, H, W) 块网格；该工具可渲染每帧 PNG 和动画 GIF。
资源感知 默认 8 GiB 保护防止 RAM/VRAM 溢出；您可以调整 max_capture_bytes 并选择要记录的时间步。
CLI 演示与验证 demo/run_attention_demo.py 可为任何支持的模型运行快速示例；测试套件（pytest）和审计脚本验证覆盖率和与原始生成的并行性。
安装 单个 pip 安装：pip install attention-map-diffusers==1.0.0。
许可与引用 MIT 许可；提供 DOI 和 BibTeX 条目用于学术用途。

如何开始（来自 README 的快速启动代码片段）

图像示例（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,
)

视频示例（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

# 将模块移回 CPU 以释放 GPU 内存
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),
)

磁盘上得到的内容

该包会写入如下目录结构：

attention/
  metadata.json                # 捕获设置及来源信息
  raw/<relation>/*.pt          # 如果使用了 --save-raw，则保存原始张量
  visuals/<relation>/aggregate/
    maps/*.png                 # 每帧或每图的热图
    maps/*.gif                 # 视频关系的动画 GIF
    overlays/…                 # 原始媒体叠加注意力图

元数据包含模型检查点、提示、捕获的时间步以及应用的任何资源限制。

谁可能使用它

  • 研究人员：探究扩散模型如何关注提示标记或图像块。
  • 提示工程师：想查看特定短语如何影响输出区域。
  • 教育者：演示生成式 AI 中交叉注意力的内部工作原理。
  • 开发者：构建 Diffusers 管道的调试工具或可视化可解释性扩展。
  • 上述所有信息均直接取自仓库的 README；未推断任何额外功能。*

相关