Initialize pretrained TAESD on the given device from the given checkpoints.
(self, encoder_path="taesd_encoder.pth", decoder_path="taesd_decoder.pth", latent_channels=None)
| 51 | latent_shift = 0.5 |
| 52 | |
| 53 | def __init__(self, encoder_path="taesd_encoder.pth", decoder_path="taesd_decoder.pth", latent_channels=None): |
| 54 | """Initialize pretrained TAESD on the given device from the given checkpoints.""" |
| 55 | super().__init__() |
| 56 | if latent_channels is None: |
| 57 | latent_channels = self.guess_latent_channels(str(encoder_path)) |
| 58 | self.encoder = Encoder(latent_channels) |
| 59 | self.decoder = Decoder(latent_channels) |
| 60 | if encoder_path is not None: |
| 61 | self.encoder.load_state_dict(torch.load(encoder_path, map_location="cpu", weights_only=True)) |
| 62 | if decoder_path is not None: |
| 63 | self.decoder.load_state_dict(torch.load(decoder_path, map_location="cpu", weights_only=True)) |
| 64 | |
| 65 | @torch.no_grad() |
| 66 | def encode(self, x): |
nothing calls this directly
no test coverage detected