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；未推斷任何額外功能。*

相關