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)