(self, in_channels, n_heads, d_head,
depth=1, dropout=0., context_dim=None,
disable_self_attn=False, use_linear=False,
use_checkpoint=True)
| 285 | NEW: use_linear for more efficiency instead of the 1x1 convs |
| 286 | """ |
| 287 | def __init__(self, in_channels, n_heads, d_head, |
| 288 | depth=1, dropout=0., context_dim=None, |
| 289 | disable_self_attn=False, use_linear=False, |
| 290 | use_checkpoint=True): |
| 291 | super().__init__() |
| 292 | if exists(context_dim) and not isinstance(context_dim, list): |
| 293 | context_dim = [context_dim] |
| 294 | self.in_channels = in_channels |
| 295 | inner_dim = n_heads * d_head |
| 296 | self.norm = Normalize(in_channels) |
| 297 | if not use_linear: |
| 298 | self.proj_in = nn.Conv2d(in_channels, |
| 299 | inner_dim, |
| 300 | kernel_size=1, |
| 301 | stride=1, |
| 302 | padding=0) |
| 303 | else: |
| 304 | self.proj_in = nn.Linear(in_channels, inner_dim) |
| 305 | |
| 306 | self.transformer_blocks = nn.ModuleList( |
| 307 | [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim[d], |
| 308 | disable_self_attn=disable_self_attn, checkpoint=use_checkpoint) |
| 309 | for d in range(depth)] |
| 310 | ) |
| 311 | if not use_linear: |
| 312 | self.proj_out = zero_module(nn.Conv2d(inner_dim, |
| 313 | in_channels, |
| 314 | kernel_size=1, |
| 315 | stride=1, |
| 316 | padding=0)) |
| 317 | else: |
| 318 | self.proj_out = zero_module(nn.Linear(in_channels, inner_dim)) |
| 319 | self.use_linear = use_linear |
| 320 | |
| 321 | def forward(self, x, context=None): |
| 322 | # note: if no context is given, cross-attention defaults to self-attention |
nothing calls this directly
no test coverage detected