| 355 | nn.init.zeros_(self.out_proj.bias) |
| 356 | |
| 357 | def forward(self, input, cond): |
| 358 | n, c, h, w = input.shape |
| 359 | q = self.q_proj(self.norm_dec(input, cond)) |
| 360 | q = q.view([n, self.n_head, c // self.n_head, h * w]).transpose(2, 3) |
| 361 | kv = self.kv_proj(self.norm_enc(cond[self.cond_key])) |
| 362 | kv = kv.view([n, -1, self.n_head * 2, c // self.n_head]).transpose(1, 2) |
| 363 | k, v = kv.chunk(2, dim=1) |
| 364 | attn_mask = (cond[self.cond_key_padding][:, None, None, :]) * -10000 |
| 365 | y = F.scaled_dot_product_attention(q, k, v, attn_mask, dropout_p=self.dropout.p) |
| 366 | y = y.transpose(2, 3).contiguous().view([n, c, h, w]) |
| 367 | return input + self.out_proj(y) |
| 368 | |
| 369 | |
| 370 | # Downsampling/upsampling |