| 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): |
| 446 | self.input_res = input_res |
| 447 | self.patch_size = patch_size |
| 448 | self.width = width |
| 449 | token_len = (input_res[0] // patch_size) * (input_res[1] // patch_size) |
| 450 | super().__init__(width, layers, heads, window_size, token_len, ConditionalResAttBlock, drop_path_rate=drop_path_rate) |
| 451 | self.positional_embedding = nn.Parameter(torch.zeros(1, token_len*2, width)) |
| 452 | nn.init.trunc_normal_(self.positional_embedding, std=0.02) |
| 453 | self.cls_embedding = nn.Parameter(torch.zeros(1, 2, width)) |
| 454 | nn.init.trunc_normal_(self.cls_embedding, std=0.02) |
| 455 | self.positional_encoding = SinusoidalPositionalEncoding(max_len=condition_len, d_model=width) |
| 456 | |
| 457 | if condition_dim is not None: |
| 458 | self.condition_proj = nn.Linear(condition_dim, width, bias=False) |
| 459 | else: |
| 460 | self.condition_proj = nn.Identity() |
| 461 | |
| 462 | self.out_proj = nn.Identity() |
| 463 | |
| 464 | self.dropout = nn.Dropout(drop_path_rate) |
| 465 | self.condition_ln = LayerNorm(width) |
| 466 | |
| 467 | def forward(self, latent, condition): |
| 468 | b, v = latent.shape[:2] |
| 469 | latent = rearrange(latent, 'b v n d -> (b v) n d') # [B, 2*N, D] |
| 470 | condition = rearrange(condition, 'b v n d -> (b v) n d') # [B, 2*N, D] |
| 471 | condition = rearrange(condition, 'b (p n) d -> (b p) n d', p=2) # [B*2, N, D] |
| 472 | |
| 473 | condition = self.condition_proj(condition) |
| 474 | latent = latent + self.positional_embedding |
| 475 | cls_embedding = self.cls_embedding.repeat_interleave(latent.shape[1]//2, dim=1).contiguous() # [1, N, D] |
| 476 | latent = latent + cls_embedding |
| 477 | |
| 478 | condition = condition + self.positional_encoding(condition).to(condition.dtype) # [B*2, N, D] |
| 479 | condition = self.condition_ln(condition) |
| 480 | condition = self.dropout(condition) |
| 481 | |
| 482 | x = super().forward(latent, condition) |
| 483 | x = self.out_proj(x) |
| 484 | |
| 485 | x = rearrange(x, '(b v) n d -> b v n d', v=v) |
| 486 | return x |
| 487 | |
| 488 | class TransformerDecoder(TransformerBase): |
| 489 | def __init__(self, token_len, width, layers, heads, window_size, encoder_dim=None): |