| 486 | return x |
| 487 | |
| 488 | class TransformerDecoder(TransformerBase): |
| 489 | def __init__(self, token_len, width, layers, heads, window_size, encoder_dim=None): |
| 490 | self.width = width |
| 491 | super().__init__(width, layers, heads, window_size, token_len, ResAttBlock) |
| 492 | self.positional_embedding = nn.Parameter(torch.zeros(1, token_len, width)) |
| 493 | nn.init.trunc_normal_(self.positional_embedding, std=0.02) |
| 494 | |
| 495 | if encoder_dim is not None and encoder_dim != width: |
| 496 | self.encoder_proj = nn.Linear(encoder_dim, width) |
| 497 | self.out_proj = nn.Linear(width, encoder_dim) |
| 498 | else: |
| 499 | self.encoder_proj = nn.Identity() |
| 500 | self.out_proj = nn.Identity() |
| 501 | |
| 502 | |
| 503 | def forward(self, latent, condition=None, reverse=False): |
| 504 | _, v = latent.shape[:2] |
| 505 | latent = rearrange(latent, 'b v n d -> (b v) n d') |
| 506 | |
| 507 | latent = self.encoder_proj(latent) |
| 508 | |
| 509 | latent = latent + self.positional_embedding.to(latent.dtype) |
| 510 | x = super().forward(latent, condition) |
| 511 | x = self.out_proj(x) |
| 512 | |
| 513 | x = rearrange(x, '(b v) n d -> b v n d', v=v) |
| 514 | return x |
| 515 | |
| 516 | class PSUpsamplerBlock(nn.Module): |
| 517 | """ |
nothing calls this directly
no outgoing calls
no test coverage detected