| 339 | |
| 340 | class CrossAttention2d(ConditionedModule): |
| 341 | def __init__(self, c_dec, c_enc, n_head, norm_dec, dropout_rate=0., |
| 342 | cond_key='cross', cond_key_padding='cross_padding'): |
| 343 | super().__init__() |
| 344 | assert c_dec % n_head == 0 |
| 345 | self.cond_key = cond_key |
| 346 | self.cond_key_padding = cond_key_padding |
| 347 | self.norm_enc = nn.LayerNorm(c_enc) |
| 348 | self.norm_dec = norm_dec(c_dec) |
| 349 | self.n_head = n_head |
| 350 | self.q_proj = nn.Conv2d(c_dec, c_dec, 1) |
| 351 | self.kv_proj = nn.Linear(c_enc, c_dec * 2) |
| 352 | self.out_proj = nn.Conv2d(c_dec, c_dec, 1) |
| 353 | self.dropout = nn.Dropout(dropout_rate) |
| 354 | nn.init.zeros_(self.out_proj.weight) |
| 355 | nn.init.zeros_(self.out_proj.bias) |
| 356 | |
| 357 | def forward(self, input, cond): |
| 358 | n, c, h, w = input.shape |