(self, latent, condition=None, reverse=False)
| 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 | """ |