wooyeolbaek/attention-map-diffusers

attention map tools for huggingface/diffusers

무엇인가요

attention‑map‑diffusers 는 Hugging Face Diffusers 라이브러리에서 실행되는 확산 모델의 내부 어텐션 텐서를 캡처하고 시각화할 수 있는 Python 패키지입니다. 이미지 생성 파이프라인(예: 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에서 직접 가져온 것이며, 추가 기능은 추측되지 않았습니다.*

관련