PQGAN: Product-Quantised Image Representation for High-Quality Image Synthesis

Pretrained PQGAN autoencoders from the ICLR 2026 paper by Denis Zavadski, Nikita Philip Tatsch and Carsten Rother.

Paper (OpenReview) · arXiv · Code

PQGAN replaces the single codebook of VQGAN with product quantisation. Each latent vector is split into S subspaces, and each subspace is quantised with its own codebook of K entries.

Models

Both models were trained on ImageNet at 256×256. Metrics are on the ImageNet validation set (50k images).

File F d K S Latent (256×256 input) PSNR ↑ rFID ↓ CMMD ↓ LPIPS ↓
checkpoints/PQGAN_F8_K512_Z128_S64 8 128 512 64 128×32×32 37.4 0.036 0.011 0.0024
checkpoints/PQGAN_F16_K128_Z128_S64 16 128 128 64 128×16×16 28.3 0.41 0.094 0.0304

F: downsampling factor, d: latent channels, K: codebook size per subspace, S: number of subspaces.

Each model has a .safetensors file with the autoencoder and quantiser weights, and a .yaml file with the model config. Discriminator and optimizer states are not included.

Usage

Install the code from GitHub (Python 3.11):

git clone https://github.com/vislearn/PQGAN.git
cd PQGAN
pip install -e .

Load a model and reconstruct an image:

import torch
import yaml
from huggingface_hub import hf_hub_download
from PIL import Image
from safetensors.torch import load_file
from torchvision import transforms

from vv.models.latent_diffusion_copy import VQModel

name = "PQGAN_F8_K512_Z128_S64"  # or "PQGAN_F16_K128_Z128_S64"
config = yaml.safe_load(open(hf_hub_download("CVL-Heidelberg/PQGAN", f"checkpoints/{name}.yaml")))
model = VQModel(**config["model"]["init_args"])
model.load_state_dict(load_file(hf_hub_download("CVL-Heidelberg/PQGAN", f"checkpoints/{name}.safetensors")))
model = model.eval().cuda()

# Models expect 256x256 inputs with ImageNet normalisation
mean, std = torch.tensor([0.485, 0.456, 0.406]), torch.tensor([0.229, 0.224, 0.225])
preprocess = transforms.Compose([transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize(mean, std)])
x = preprocess(Image.open("image.jpg").convert("RGB")).unsqueeze(0).cuda()

with torch.no_grad():
    quant, _, (_, _, indices) = model.encode(x)  # quant: (1, 128, 32, 32) for F=8, (1, 128, 16, 16) for F=16
    rec = model.decode(quant)

rec = (rec.cpu() * std[:, None, None] + mean[:, None, None]).clamp(0, 1)  # back to [0, 1]

indices is a list with one tensor of codebook indices per subspace (64 entries).

Citation

@inproceedings{zavadski2026pqgan,
  title     = {{PQGAN}: Product-Quantised Image Representation for High-Quality Image Synthesis},
  author    = {Zavadski, Denis and Tatsch, Nikita Philip and Rother, Carsten},
  booktitle = {The Fourteenth International Conference on Learning Representations},
  year      = {2026},
  url       = {https://openreview.net/forum?id=D8oqcochgq}
}
Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Dataset used to train CVL-Heidelberg/PQGAN

Paper for CVL-Heidelberg/PQGAN