Views
No views yet




1# build DC-AE models
2# full DC-AE model list: https://huggingface.co/collections/mit-han-lab/dc-ae-670085b9400ad7197bb1009b
3from efficientvit.ae_model_zoo import DCAE_HF
4
5dc_ae = DCAE_HF.from_pretrained(f"mit-han-lab/dc-ae-f64c128-in-1.0")
6
7# encode
8from PIL import Image
9import torch
10import torchvision.transforms as transforms
11from torchvision.utils import save_image
12from efficientvit.apps.utils.image import DMCrop
13
14device = torch.device("cuda")
15dc_ae = dc_ae.to(device).eval()
16
17transform = transforms.Compose([
18 DMCrop(512), # resolution
19 transforms.ToTensor(),
20 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
21])
22image = Image.open("assets/fig/girl.png")
23x = transform(image)[None].to(device)
24latent = dc_ae.encode(x)
25print(latent.shape)
26
27# decode
28y = dc_ae.decode(latent)
29save_image(y * 0.5 + 0.5, "demo_dc_ae.png")1# build DC-AE-Diffusion models
2# full DC-AE-Diffusion model list: https://huggingface.co/collections/mit-han-lab/dc-ae-diffusion-670dbb8d6b6914cf24c1a49d
3from efficientvit.diffusion_model_zoo import DCAE_Diffusion_HF
4
5dc_ae_diffusion = DCAE_Diffusion_HF.from_pretrained(f"mit-han-lab/dc-ae-f64c128-in-1.0-uvit-h-in-512px-train2000k")
6
7# denoising on the latent space
8import torch
9import numpy as np
10from torchvision.utils import save_image
11
12torch.set_grad_enabled(False)
13device = torch.device("cuda")
14dc_ae_diffusion = dc_ae_diffusion.to(device).eval()
15
16seed = 0
17torch.manual_seed(seed)
18torch.cuda.manual_seed_all(seed)
19eval_generator = torch.Generator(device=device)
20eval_generator.manual_seed(seed)
21
22prompts = torch.tensor(
23 [279, 333, 979, 936, 933, 145, 497, 1, 248, 360, 793, 12, 387, 437, 938, 978], dtype=torch.int, device=device
24)
25num_samples = prompts.shape[0]
26prompts_null = 1000 * torch.ones((num_samples,), dtype=torch.int, device=device)
27latent_samples = dc_ae_diffusion.diffusion_model.generate(prompts, prompts_null, 6.0, eval_generator)
28latent_samples = latent_samples / dc_ae_diffusion.scaling_factor
29
30# decode
31image_samples = dc_ae_diffusion.autoencoder.decode(latent_samples)
32save_image(image_samples * 0.5 + 0.5, "demo_dc_ae_diffusion.png", nrow=int(np.sqrt(num_samples)))@article{chen2024deep,
title={Deep Compression Autoencoder for Efficient High-Resolution Diffusion Models},
author={Chen, Junyu and Cai, Han and Chen, Junsong and Xie, Enze and Yang, Shang and Tang, Haotian and Li, Muyang and Lu, Yao and Han, Song},
journal={arXiv preprint arXiv:2410.10733},
year={2024}
}