docs/comp_attn_design.md

# Comp-Attn 设计文档

> 基于论文 "Comp-Attn: Present-and-Align Attention for Compositional Video Generation" 的实现总结

## 概述

Comp-Attn 是一种 **composition-aware cross-attention** 变体，采用 **"Present-and-Align"** 范式，用于解决多主体视频生成中的两个核心问题：

| 挑战 | 描述 | 解决方案 |
|-----|------|---------|
| **Subject Presence** | 并非所有主体都能在视频中呈现 | SCI（条件层面） |
| **Inter-subject Relations** | 主体间的交互和空间关系错位 | LAM（注意力层面） |

## 核心组件

### 1. Subject-aware Condition Interpolation (SCI)

SCI 在 **条件编码阶段** 增强每个主体的语义表示，确保所有主体都能被"召回"。

#### 工作原理

```
┌─────────────────────────────────────────────────────────────┐
│  输入: prompt 中的主体 tokens + 独立编码的 anchor embeddings │
├─────────────────────────────────────────────────────────────┤
│  Step 1: 计算语义显著性 (Saliency)                           │
│          - 比较 prompt 中主体 token 与 anchor 的余弦相似度    │
│          - 使用 softmax(τ=0.2) 归一化得到显著性分数            │
│          - 低显著性 = 主体在 prompt 上下文中被"淹没"          │
├─────────────────────────────────────────────────────────────┤
│  Step 2: 计算语义差异 (Delta)                                │
│          - delta_i = anchor_i * N - Σ(anchors)              │
│          - 表示每个主体相对于其他主体的独特语义               │
├─────────────────────────────────────────────────────────────┤
│  Step 3: 自适应插值                                          │
│          - scale = ω * (1 - saliency)                       │
│          - context' = context + delta * scale               │
│          - 显著性越低，增强越多；早期时间步增强越强           │
└─────────────────────────────────────────────────────────────┘
```

#### 关键代码

```python
# 显著性计算
def compute_saliency(prompt_vecs, anchor_vecs, tau=0.2):
    cosine = cosine_similarity(prompt_vecs, anchor_vecs)
    scores = exp(cosine / tau)
    return scores.diagonal() / scores.sum(dim=1)

# 应用 SCI
omega = 1.0 - (timestep / 1000.0)  # 时间步调度
scale = omega * (1.0 - saliency)
context = context + delta * scale
```

### 2. Layout-forcing Attention Modulation (LAM)

LAM 在 **注意力计算阶段** 动态调制注意力分布，使其与预定义的空间布局对齐。

#### 工作原理

```
┌─────────────────────────────────────────────────────────────┐
│  输入: Q (视频 tokens), K/V (文本 tokens), Layout (bbox)     │
├─────────────────────────────────────────────────────────────┤
│  Step 1: 计算原始 attention scores                           │
│          attn_scores = Q @ K^T / sqrt(d)                    │
├─────────────────────────────────────────────────────────────┤
│  Step 2: 构建调制函数                                        │
│          g_plus  = max(scores) - scores  (增强调制)         │
│          g_minus = min(scores) - scores  (抑制调制)         │
├─────────────────────────────────────────────────────────────┤
│  Step 3: 计算 IOU 引导的强度                                 │
│          - adapt_mask: 当前注意力分布 > mean 的区域          │
│          - layout_mask: 目标 bbox 区域                       │
│          - iou = intersection / union                       │
│          - strength = 1 - iou (IOU 越低，调制越强)          │
├─────────────────────────────────────────────────────────────┤
│  Step 4: 应用动态调制                                        │
│          - 在 bbox 内增强对应主体 token (g_plus)             │
│          - 抑制其他主体 token (g_minus)                      │
│          - bias = g_k * strength * layout_mask              │
│          - final_scores = attn_scores + bias                │
└─────────────────────────────────────────────────────────────┘
```

#### LAM vs 传统 Layout Control

| 方法 | 机制 | 缺点 |
|-----|------|-----|
| **硬掩码 (Hard Mask)** | 强制注意力只在 bbox 内 | 无法适应多样的物体形状 |
| **LAM (Ours)** | IOU 引导的动态软调制 | ✅ 灵活适应不同形状 |

### 3. 关键帧插值

遵循论文附录 F 的设计，支持用 4 个关键帧描述运动轨迹，然后线性插值到所有帧：

```python
# 输入: 4 个关键帧的 bbox
keyframe_bboxes = [frame_0, frame_1, frame_2, frame_3]

# 输出: 81 帧的 bbox（线性插值）
all_bboxes = F.interpolate(keyframe_bboxes, size=81, mode='linear')
```

## 架构设计

```
┌─────────────────────────────────────────────────────────────┐
│                    WanVideoCompAttnPipeline                 │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  ┌──────────────────┐    ┌──────────────────┐               │
│  │  CompAttnUnit    │───▶│  CompAttnMerge   │               │
│  │  (SCI 预处理)    │    │  (CFG 合并)      │               │
│  └──────────────────┘    └──────────────────┘               │
│           │                                                 │
│           ▼                                                 │
│  ┌──────────────────────────────────────────────────────┐  │
│  │  model_fn_wrapper (动态注入)                          │  │
│  │  - apply_sci(): 修改 context                         │  │
│  │  - build_layout_mask(): 构建空间掩码                  │  │
│  └──────────────────────────────────────────────────────┘  │
│           │                                                 │
│           ▼                                                 │
│  ┌──────────────────────────────────────────────────────┐  │
│  │  patched cross_attention                              │  │
│  │  - lam_attention(): IOU 引导的注意力调制              │  │
│  └──────────────────────────────────────────────────────┘  │
│                                                             │
└─────────────────────────────────────────────────────────────┘
```

## Prompt-BBox 绑定机制

### 使用变量拼接（推荐）

使用 Python 变量定义主体，通过 f-string 拼接 prompt，代码层面天然表达绑定关系：

```python
# 定义主体变量
subject0 = "red car"
subject1 = "blue bicycle"

# 使用变量拼接 prompt
prompt = f"A {subject0} drives left, a {subject1} rides right"

# subjects 和 bboxes 顺序一一对应
subjects = [subject0, subject1]
bboxes = [bbox0, bbox1]  # bbox0 对应 subject0, bbox1 对应 subject1
```

### 绑定原理

```
┌─────────────────────────────────────────────────────────────────────────┐
│  1. 用户定义 subjects 和 bboxes（顺序一一对应）                          │
│                                                                         │
│     subject0 = "red car"                                                │
│     subject1 = "blue bicycle"                                           │
│     subjects = [subject0, subject1]                                     │
│     bboxes = [bbox0, bbox1]                                             │
├─────────────────────────────────────────────────────────────────────────┤
│  2. Tokenize prompt，搜索每个 subject 的 token 位置                      │
│                                                                         │
│     prompt = f"A {subject0} drives while a {subject1} rides..."        │
│                  ↑↑↑↑↑↑↑                  ↑↑↑↑↑↑↑↑↑↑↑↑                 │
│                  token indices            token indices                 │
├─────────────────────────────────────────────────────────────────────────┤
│  3. 建立 subject_token_mask (关联 token 位置与 bbox)                     │
│                                                                         │
│     subject_token_mask[0, ...] = True  # subject0 tokens                │
│     subject_token_mask[1, ...] = True  # subject1 tokens                │
├─────────────────────────────────────────────────────────────────────────┤
│  4. 在推理时应用：                                                       │
│     - SCI: 根据 mask 增强对应 token 的语义                               │
│     - LAM: 在 bbox 区域内增强对应 token 的注意力                          │
└─────────────────────────────────────────────────────────────────────────┘
```

### 重要约束

1. **subjects 必须在 prompt 中出现**：使用变量拼接确保这一点
2. **顺序一一对应**：`bboxes[i]` 对应 `subjects[i]`
3. **精确匹配**：subject 字符串需要能被分词器识别为完整的 token 序列

## 使用方法

### 基本用法

```python
from diffsynth.pipelines.wan_video_comp_attn import WanVideoCompAttnPipeline
from diffsynth.models.comp_attn_model import CompAttnConfig

# 创建 pipeline
pipe = WanVideoCompAttnPipeline.from_pretrained(...)

# 1. 定义主体变量
subject0 = "red car"
subject1 = "blue bicycle"

# 2. 定义运动轨迹 (4 个关键帧)
bbox0 = [(100, 250, 220, 380), (200, 250, 320, 380), (300, 250, 420, 380), (400, 250, 520, 380)]
bbox1 = [(350, 260, 410, 400), (450, 260, 510, 400), (550, 260, 610, 400), (650, 260, 710, 400)]

# 3. 使用变量拼接 prompt
prompt = f"A {subject0} drives from left to center while a {subject1} rides to the right"

# 4. 配置 Comp-Attn（顺序一一对应）
comp_attn = CompAttnConfig(
    subjects=[subject0, subject1],  # 变量列表
    bboxes=[bbox0, bbox1],          # 对应的 bbox 列表
    enable_sci=True,
    enable_lam=True,
    interpolate=True,
)

# 5. 生成视频
video = pipe(prompt=prompt, comp_attn=comp_attn)
```

### 运动轨迹辅助函数

```python
def create_moving_bbox(start_x, end_x, y_center, box_width, box_height, num_keyframes=4):
    """创建从 start_x 移动到 end_x 的关键帧 bbox 序列"""
    keyframes = []
    for i in range(num_keyframes):
        progress = i / (num_keyframes - 1)
        center_x = start_x + (end_x - start_x) * progress
        left = center_x - box_width / 2
        right = center_x + box_width / 2
        top = y_center - box_height / 2
        bottom = y_center + box_height / 2
        keyframes.append((left, top, right, bottom))
    return keyframes

# 使用示例
car_trajectory = create_moving_bbox(
    start_x=100, end_x=500,
    y_center=300,
    box_width=120, box_height=80,
)
```

### 配置参数

| 参数 | 类型 | 默认值 | 说明 |
|-----|------|-------|------|
| `subjects` | List[str] | - | 主体名称列表（必须出现在 prompt 中） |
| `bboxes` | List | None | 每个主体每个关键帧的 bbox |
| `enable_sci` | bool | True | 是否启用 SCI |
| `enable_lam` | bool | True | 是否启用 LAM |
| `temperature` | float | 0.2 | 显著性计算的温度参数 τ |
| `apply_to_negative` | bool | False | 是否对负样本 prompt 应用 |
| `interpolate` | bool | False | 是否对关键帧 bbox 进行插值 |

### BBox 格式

```python
# 格式: (left, top, right, bottom) 像素坐标
# 支持多种输入形式:

# 1. 静态布局（所有帧相同）
bbox = (100, 200, 300, 400)

# 2. 关键帧布局（4帧，会被插值）
bboxes = [
    (100, 200, 300, 400),  # 关键帧 0
    (150, 200, 350, 400),  # 关键帧 1
    (200, 200, 400, 400),  # 关键帧 2
    (250, 200, 450, 400),  # 关键帧 3
]
```

## 与论文的对应关系

| 论文章节 | 实现位置 | 状态 |
|---------|---------|-----|
| Sec 3.2 SCI | `compute_saliency()`, `compute_delta()`, `apply_sci()` | ✅ |
| Sec 3.3 LAM | `lam_attention()`, `build_layout_mask_from_bboxes()` | ✅ |
| Appendix D 关键帧插值 | `interpolate_bboxes()` | ✅ |
| Training-free 集成 | `patch_cross_attention()`, `wrap_model_fn()` | ✅ |

## 性能特点

根据论文数据：

- **T2V-CompBench** 性能提升: +15.7% (Wan2.1-14B), +11.7% (Wan2.2-A14B)
- **推理时间增加**: 仅 ~5%
- **兼容性**: Wan, CogVideoX, VideoCrafter2, FLUX

## 注意事项

1. **主体名称必须精确匹配**: `subjects` 中的字符串必须能在 prompt 中找到对应的 token
2. **BBox 使用像素坐标**: 不是归一化坐标
3. **关键帧数量**: 推荐使用 4 个关键帧描述运动轨迹
4. **温度参数**: τ 过小会导致显著性估计不稳定，过大会削弱增强效果
5. **State tokens**: per-frame state control adds extra context tokens and attention bias. Keep the number of states small to reduce overhead.

## Per-frame State Control

You can inject per-frame instance states (e.g., "running" -> "idle") with:

- `state_texts`: list of state names per subject
- `state_weights`: per-frame weights `(M, F, S)` or `(B, M, F, S)`
- `state_scale`: bias strength for state tokens
- `state_template`: default `"{subject} is {state}"`

The implementation appends state tokens to the context and applies a per-frame attention bias
based on the current token time index.

## 参考

- 论文: "Comp-Attn: Present-and-Align Attention for Compositional Video Generation"
- 作者: Hongyu Zhang, Yufan Deng, et al. (Peking University, Tsinghua University)