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