| 182 | Note: this is a single-head self-attention operation |
| 183 | """ |
| 184 | def __init__(self, in_channels, natten_kernel_size=-1, use_null_attention=False): |
| 185 | super().__init__() |
| 186 | self.in_channels = in_channels |
| 187 | |
| 188 | self.norm = Normalize(in_channels) |
| 189 | conv_kwargs = dict(kernel_size=1, stride=1, padding=0) |
| 190 | self.q = nn.Conv2d(in_channels, in_channels, **conv_kwargs) |
| 191 | self.k = nn.Conv2d(in_channels, in_channels, **conv_kwargs) |
| 192 | self.v = nn.Conv2d(in_channels, in_channels, **conv_kwargs) |
| 193 | self.proj_out = nn.Conv2d(in_channels, in_channels, **conv_kwargs) |
| 194 | |
| 195 | if natten_kernel_size > -1: |
| 196 | assert NATTEN_IS_AVAILBLE, "natten_kernel_size > -1 but natten is not available" |
| 197 | assert (natten_kernel_size % 2) == 1, 'natten_kernel_size must be odd' |
| 198 | self.natten_kernel_size = natten_kernel_size |
| 199 | self.use_null_attention = use_null_attention |
| 200 | |
| 201 | def forward(self, x): |
| 202 | h_ = x |