| 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): |