| 420 | return x |
| 421 | |
| 422 | class TransformerEncoder(TransformerBase): |
| 423 | def __init__(self, input_res, in_channels, patch_size, width, layers, heads, window_size): |
| 424 | self.input_res = input_res |
| 425 | self.patch_size = patch_size |
| 426 | token_len = (self.input_res[0] // patch_size) * (self.input_res[1] // patch_size) |
| 427 | super().__init__(width, layers, heads, window_size, token_len, ResAttBlock) |
| 428 | self.conv = nn.Conv2d(in_channels=in_channels, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False) |
| 429 | self.positional_encoding = SinusoidalPositionalEncoding(max_len=token_len, d_model=width) |
| 430 | |
| 431 | def forward(self, x, condition=None): |
| 432 | _, v = x.shape[:2] |
| 433 | x = rearrange(x, 'b v c h w -> (b v) c h w') |
| 434 | x = self.conv(x) |
| 435 | x = x.reshape(x.shape[0], x.shape[1], -1) |
| 436 | x = x.permute(0, 2, 1) |
| 437 | |
| 438 | x = x + self.positional_encoding(x).to(x.dtype) |
| 439 | x = super().forward(x, condition) |
| 440 | |
| 441 | x = rearrange(x, 'b (v n) d -> b v n d', v=v) |
| 442 | return x |
| 443 | |
| 444 | class TransformerConditionalDecoder(TransformerBase): |
| 445 | def __init__(self, input_res, patch_size, width, layers, heads, window_size, encoder_dim=None, condition_len=576, condition_dim=None, drop_path_rate=0.1): |