| 461 | |
| 462 | |
| 463 | class AutoencoderKL(nn.Module): |
| 464 | def __init__( |
| 465 | self, |
| 466 | ckpt_path: str = None, |
| 467 | ddconfig=DEFAULT_DDCONFIG, |
| 468 | embed_dim: int = 4, |
| 469 | scale: float = 0.18215, # SD: 0.18215 |
| 470 | shift: float = 0.0 |
| 471 | ): |
| 472 | super().__init__() |
| 473 | self.encoder = Encoder(**ddconfig) |
| 474 | self.decoder = Decoder(**ddconfig) |
| 475 | assert ddconfig["double_z"] |
| 476 | self.quant_conv = nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1) |
| 477 | self.post_quant_conv = nn.Conv2d(embed_dim, ddconfig["z_channels"], 1) |
| 478 | self.embed_dim = embed_dim |
| 479 | |
| 480 | self.scale = scale |
| 481 | self.shift = shift |
| 482 | |
| 483 | if exists(ckpt_path): |
| 484 | assert os.path.exists(ckpt_path), f'[AutoencoderKL] Checkpoint {ckpt_path} not found!' |
| 485 | print(f'[AutoencoderKL] Loading checkpoint from {ckpt_path}') |
| 486 | if torch.cuda.is_available(): |
| 487 | self.load_state_dict(torch.load(ckpt_path, weights_only=True)) |
| 488 | else: |
| 489 | self.load_state_dict(torch.load(ckpt_path, weights_only=True, map_location=torch.device('cpu'))) |
| 490 | else: |
| 491 | import warnings |
| 492 | warnings.warn(f'[AutoencoderKL] No checkpoint provided. Random initialization.') |
| 493 | |
| 494 | @torch.no_grad() |
| 495 | def encode(self, x: torch.Tensor, return_posterior=False): |
| 496 | """ |
| 497 | Args: |
| 498 | x: input tensor (B, C, H, W) in range [-1, 1] scaled with |
| 499 | self.scale and shifted with self.shift |
| 500 | return_posterior: return the posterior distribution |
| 501 | """ |
| 502 | h = self.encoder(x) |
| 503 | moments = self.quant_conv(h) |
| 504 | posterior = DiagonalGaussianDistribution(moments) |
| 505 | if return_posterior: |
| 506 | return posterior |
| 507 | latent = posterior.mode() |
| 508 | return (latent + self.shift) * self.scale |
| 509 | |
| 510 | @torch.no_grad() |
| 511 | def decode(self, z: torch.Tensor): |
| 512 | """ |
| 513 | Args: |
| 514 | z: latent code tensor (B, C, H, W) |
| 515 | """ |
| 516 | z = z / self.scale + self.shift |
| 517 | z = self.post_quant_conv(z) |
| 518 | dec = self.decoder(z) |
| 519 | return dec |
| 520 | |