| import torch | |
| from models import VQVAE, build_vae_var | |
| from dataset.imagenet_dataset import get_train_transforms | |
| from PIL import Image | |
| from torchvision import transforms | |
| device = 'mps' | |
| patch_nums = (1, 2, 3, 4, 5, 6, 8, 10, 13, 16) | |
| vae, var = build_vae_var( | |
| V=4096, Cvae=32, ch=160, share_quant_resi=4, | |
| device=device, patch_nums=patch_nums, | |
| num_classes=1000, depth=16, shared_aln=False, | |
| ) | |
| var_ckpt='var_d16.pth' | |
| vae_ckpt='vae_ch160v4096z32.pth' | |
| var.load_state_dict(torch.load(var_ckpt, map_location=device), strict=True) | |
| vae.load_state_dict(torch.load(vae_ckpt, map_location=device), strict=True) |