(self, token, device, path)
| 380 | |
| 381 | class TorchVAEEncoder(torch.nn.Module): |
| 382 | def __init__(self, token, device, path): |
| 383 | super().__init__() |
| 384 | self.path = path |
| 385 | self.vae_encoder = AutoencoderKL.from_pretrained(self.path, subfolder="vae", use_auth_token=token).to(device) |
| 386 | |
| 387 | def forward(self, x): |
| 388 | return self.vae_encoder.encode(x).latent_dist.sample() |