(self, x)
| 55 | self.attention = QKVAttention(self.num_heads) |
| 56 | |
| 57 | def forward(self, x): |
| 58 | b, c, *_spatial = x.shape |
| 59 | x = x.reshape(b, c, -1) # NC(HW) |
| 60 | x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1) # NC(HW+1) |
| 61 | x = x + self.positional_embedding[None, :, :].to(x.dtype) # NC(HW+1) |
| 62 | x = self.qkv_proj(x) |
| 63 | x = self.attention(x) |
| 64 | x = self.c_proj(x) |
| 65 | return x[:, :, 0] |
| 66 | |
| 67 | |
| 68 | class TimestepBlock(nn.Module): |