Views
No views yet


1from PIL import Image
2import torch
3import torchvision.transforms as transforms
4from torchvision.utils import save_image
5from dc_gen.ae_model_zoo import DCAE_HF
6
7device = torch.device("cuda")
8dc_ae_lite = DCAE_HF.from_pretrained("dc-ai/dc-ae-lite-f32c32").to(device).eval()
9
10transform = transforms.Compose([
11 transforms.CenterCrop((1024,1024)),
12 transforms.ToTensor(),
13 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
14])
15
16image = Image.open("assets/fig/girl.png")
17
18x = transform(image)[None].to(device)
19latent = dc_ae_lite.encode(x)
20print(f"latent shape: {latent.shape}")
21
22y = dc_ae_lite.decode(latent)
23save_image(y * 0.5 + 0.5, "demo_dc_ae_lite.png")