VQGAN with DCT

Usage

VQGAN-DCT:

ckpt_path = "/path/to/kl16.ckpt"
vae = VQModel(ckpt_path=ckpt_path).eval().cuda()

# Load image
image = torch.randn(10, 3, 256, 256).cuda() # (B, C, H, W)

# Encoding
encoded = vae.encode(image)
token_ids = encoded["token_ids"] # (B, H, W) - Tokenized ids - torch.int
quant = encoded["quant"] # (B, H, W, C) - Quantified discrete embeddings

# Decoding
decoded = vae.decode(token_ids) # (B, C, H, W) - Decoded image

VQGAN-Vanilla (set use_dct=False) :

ckpt_path = "/path/to/kl16-vanilla.ckpt"
vae = VQModel(use_dct=False, ckpt_path=ckpt_path).eval().cuda()

# Load image
image = torch.randn(10, 3, 256, 256).cuda() # (B, C, H, W)

# Encoding
encoded = vae.encode(image)
token_ids = encoded["token_ids"] # (B, H, W) - Tokenized ids - torch.int
quant = encoded["quant"] # (B, H, W, C) - Quantified discrete embeddings

# Decoding
decoded = vae.decode(token_ids) # (B, C, H, W) - Decoded image