| 136 | |
| 137 | class SpatialSelfAttention(nn.Module): |
| 138 | def __init__(self, in_channels): |
| 139 | super().__init__() |
| 140 | self.in_channels = in_channels |
| 141 | |
| 142 | self.norm = Normalize(in_channels) |
| 143 | self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) |
| 144 | self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) |
| 145 | self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) |
| 146 | self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) |
| 147 | |
| 148 | def forward(self, x): |
| 149 | h_ = x |