(self, x)
| 52 | self.attention = QKVAttention(self.num_heads) |
| 53 | |
| 54 | def forward(self, x): |
| 55 | b, c, *_spatial = x.shape |
| 56 | x = x.reshape(b, c, -1) # NC(HW) |
| 57 | x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1) # NC(HW+1) |
| 58 | x = x + self.positional_embedding[None, :, :].to(x.dtype) # NC(HW+1) |
| 59 | x = self.qkv_proj(x) |
| 60 | x = self.attention(x) |
| 61 | x = self.c_proj(x) |
| 62 | return x[:, :, 0] |
| 63 | |
| 64 | |
| 65 | class TimestepBlock(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected