| 130 | |
| 131 | class CrossAttn(nn.Module): |
| 132 | def __init__(self, |
| 133 | q_dim, |
| 134 | kv_dim, |
| 135 | hidden_dim, |
| 136 | num_heads, |
| 137 | out_dim=None, |
| 138 | qkv_bias=False, |
| 139 | qk_scale=None, |
| 140 | attn_drop=0., |
| 141 | proj_drop=0., |
| 142 | qkv_fuse=False): |
| 143 | super().__init__() |
| 144 | if out_dim is None: |
| 145 | out_dim = q_dim |
| 146 | self.num_heads = num_heads |
| 147 | head_dim = hidden_dim // num_heads |
| 148 | self.scale = qk_scale or head_dim**-0.5 |
| 149 | self.qkv_fuse = qkv_fuse |
| 150 | |
| 151 | self.q_proj = nn.Linear(q_dim, hidden_dim, bias=qkv_bias) |
| 152 | self.k_proj = nn.Linear(kv_dim, hidden_dim, bias=qkv_bias) |
| 153 | self.v_proj = nn.Linear(kv_dim, hidden_dim, bias=qkv_bias) |
| 154 | self.attn_drop = nn.Dropout(attn_drop) |
| 155 | self.proj = nn.Linear(hidden_dim, out_dim) |
| 156 | self.proj_drop = nn.Dropout(proj_drop) |
| 157 | |
| 158 | def forward(self, query, key, value=None, mask=None): |
| 159 | B, N, C = query.shape |