| 241 | |
| 242 | class MemoryEfficientAttnBlock(nn.Module): |
| 243 | def __init__(self, in_channels): |
| 244 | super().__init__() |
| 245 | self.in_channels = in_channels |
| 246 | |
| 247 | self.norm = Normalize(in_channels) |
| 248 | self.q = torch.nn.Conv2d(in_channels, |
| 249 | in_channels, |
| 250 | kernel_size=1, |
| 251 | stride=1, |
| 252 | padding=0) |
| 253 | self.k = torch.nn.Conv2d(in_channels, |
| 254 | in_channels, |
| 255 | kernel_size=1, |
| 256 | stride=1, |
| 257 | padding=0) |
| 258 | self.v = torch.nn.Conv2d(in_channels, |
| 259 | in_channels, |
| 260 | kernel_size=1, |
| 261 | stride=1, |
| 262 | padding=0) |
| 263 | self.proj_out = torch.nn.Conv2d(in_channels, |
| 264 | in_channels, |
| 265 | kernel_size=1, |
| 266 | stride=1, |
| 267 | padding=0) |
| 268 | self.attention_op: Optional[Any] = None |
| 269 | |
| 270 | |
| 271 | def forward(self, x): |