| 82 | |
| 83 | # Stage 1 |
| 84 | class CrossAttention(nn.Module): |
| 85 | def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None): |
| 86 | super(CrossAttention, self).__init__() |
| 87 | assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}." |
| 88 | |
| 89 | self.dim = dim |
| 90 | self.num_heads = num_heads |
| 91 | head_dim = dim // num_heads |
| 92 | self.scale = qk_scale or head_dim**-0.5 |
| 93 | self.kv1 = nn.Linear(dim, dim * 2, bias=qkv_bias) |
| 94 | self.kv2 = nn.Linear(dim, dim * 2, bias=qkv_bias) |
| 95 | |
| 96 | def forward(self, x1, x2): |
| 97 | B, N, C = x1.shape |
| 98 | q1 = x1.reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3).contiguous() |
| 99 | q2 = x2.reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3).contiguous() |
| 100 | k1, v1 = ( |
| 101 | self.kv1(x1).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).contiguous() |
| 102 | ) |
| 103 | k2, v2 = ( |
| 104 | self.kv2(x2).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).contiguous() |
| 105 | ) |
| 106 | |
| 107 | ctx1 = (k1.transpose(-2, -1) @ v1) * self.scale |
| 108 | ctx1 = ctx1.softmax(dim=-2) |
| 109 | ctx2 = (k2.transpose(-2, -1) @ v2) * self.scale |
| 110 | ctx2 = ctx2.softmax(dim=-2) |
| 111 | |
| 112 | x1 = (q1 @ ctx2).permute(0, 2, 1, 3).reshape(B, N, C).contiguous() |
| 113 | x2 = (q2 @ ctx1).permute(0, 2, 1, 3).reshape(B, N, C).contiguous() |
| 114 | |
| 115 | return x1, x2 |
| 116 | |
| 117 | |
| 118 | class CrossPath(nn.Module): |