| 448 | |
| 449 | |
| 450 | class AutoencoderKL(nn.Module): |
| 451 | def __init__(self, embed_dim, ch_mult, use_variational=True, ckpt_path=None): |
| 452 | super().__init__() |
| 453 | self.encoder = Encoder(ch_mult=ch_mult, z_channels=embed_dim) |
| 454 | self.decoder = Decoder(ch_mult=ch_mult, z_channels=embed_dim) |
| 455 | self.use_variational = use_variational |
| 456 | mult = 2 if self.use_variational else 1 |
| 457 | self.quant_conv = torch.nn.Conv2d(2 * embed_dim, mult * embed_dim, 1) |
| 458 | self.post_quant_conv = torch.nn.Conv2d(embed_dim, embed_dim, 1) |
| 459 | self.embed_dim = embed_dim |
| 460 | if ckpt_path is not None: |
| 461 | self.init_from_ckpt(ckpt_path) |
| 462 | |
| 463 | def init_from_ckpt(self, path): |
| 464 | sd = torch.load(path, map_location="cpu")["model"] |
| 465 | msg = self.load_state_dict(sd, strict=False) |
| 466 | print("Loading pre-trained KL-VAE") |
| 467 | print("Missing keys:") |
| 468 | print(msg.missing_keys) |
| 469 | print("Unexpected keys:") |
| 470 | print(msg.unexpected_keys) |
| 471 | print(f"Restored from {path}") |
| 472 | |
| 473 | def encode(self, x): |
| 474 | h = self.encoder(x) |
| 475 | moments = self.quant_conv(h) |
| 476 | if not self.use_variational: |
| 477 | moments = torch.cat((moments, torch.ones_like(moments)), 1) |
| 478 | posterior = DiagonalGaussianDistribution(moments) |
| 479 | return posterior |
| 480 | |
| 481 | def decode(self, z): |
| 482 | z = self.post_quant_conv(z) |
| 483 | dec = self.decoder(z) |
| 484 | return dec |
| 485 | |
| 486 | def forward(self, inputs, disable=True, train=True, optimizer_idx=0): |
| 487 | if train: |
| 488 | return self.training_step(inputs, disable, optimizer_idx) |
| 489 | else: |
| 490 | return self.validation_step(inputs, disable) |