This repository contains the PyTorch implementation of Parallel Rollout Approximation (PRA) for class-conditional pixel-space autoregressive image generation.
Paper: Parallel Rollout Approximation for Pixel-Space Autoregressive Image Generation
The environment file is provided in environment.yml.
conda env create -f environment.yml
conda activate pra
pip install -r requirements.txtIf the PyTorch wheel index or FlashAttention wheel URL is incompatible with your CUDA, PyTorch, Python, or CXX11 ABI setup, replace the entries in requirements.txt with matching packages from the PyTorch and FlashAttention releases.
| Name | params | FID (256x256) | weight |
|---|---|---|---|
| PRA-S | 135M | 2.58 | PRA_S.pt |
| PRA-B | 250M | 2.21 | PRA_B.pt |
| PRA-L | 511M | 1.94 | PRA_L.pt |
sample_ddp.py supports distributed class-balanced sampling. Enable at least one of --save-png or --compute-metric.
Save PNG samples:
ckpt=your_ckpt_path
sample_dir=your_result_path
torchrun --nnodes=1 --nproc_per_node=4 --node_rank=0 \
sample_ddp.py \
--ckpt $ckpt \
--sample-dir $sample_dir \
--model PRA-L \
--image-size 256 \
--patch-size 16 \
--latent-dim 16 \
--cfg-scale 4.1 \
--sample-steps 100 \
--sampler euler_maruyama \
--per-proc-batch-size 200 \
--sample-mask-rate 0.9 \
--token-mask-rate 0.5 \
--save-pngWhen --compute-metric is enabled, sample_ddp.py extracts Inception features for generated images and the reference batch, gathers features across ranks, and writes a JSON file containing FID, Inception Score, precision, and recall.
For ImageNet 256×256, use the standard guided-diffusion reference batch:
You should download the reference batch, inception weights.
VIRTUAL_imagenet256_labeled.npz
weights-inception-2015-12-05-6726825d.pth
Compute metrics:
ckpt=your_ckpt_path
sample_dir=your_result_path
torchrun --nnodes=1 --nproc_per_node=4 --node_rank=0 \
sample_ddp.py \
--ckpt $ckpt \
--sample-dir $sample_dir \
--model PRA-L \
--image-size 256 \
--patch-size 16 \
--latent-dim 16 \
--cfg-scale 4.1 \
--sample-steps 100 \
--sampler euler_maruyama \
--per-proc-batch-size 200 \
--sample-mask-rate 0.9 \
--token-mask-rate 0.5 \
--inception-weights weights-inception-2015-12-05-6726825d.pth \
--ref-batch VIRTUAL_imagenet256_labeled.npz \
--compute-metric The training code supports ImageNet either as a standard ImageFolder directory or directly from the original ImageNet training tar file.
Examples:
# Original ImageNet tar, no decompression required
data_path=/path/to/ILSVRC2012_img_train.tar
# Or an ImageFolder-style directory
data_path=/path/to/imagenet/trainImages are center/random cropped to --image-size, normalized to [-1, 1], and split into raster-ordered pixel patches.
PRA is trained end-to-end in a single stage. There is no separate VAE/tokenizer training step.
Example: train PRA-L on ImageNet 256×256 with 8×A100 GPUs:
data_path=/path/to/ILSVRC2012_img_train.tar
result_path=your_result_path
torchrun --nproc_per_node=8 --node_rank=0 \
train.py \
--results-dir $result_path \
--data-path $data_path \
--image-size 256 \
--model PRA-L \
--epochs 400 \
--patch-size 16 \
--cls-token-num 16 \
--global-batch-size 512 \
--sample-mask-rate 0.9 \
--token-mask-rate 0.5 \
--latent-noise-rate 0.9 \
--latent-noise-t 0.5 \
--pixel-weight 1.0 \
--encoder-grad-scale 0.3 \
--ema-decay 0.9999 \
--ar-gt-pixel-replace-rate 0.04 \
--noise-t-switch-epoch 350 \
--noise-t-high 0.7To train smaller models, set --model to PRA-B or PRA-S.
Note: We use torch.compile for acceleration. Occasionally the TorchInductor compile step can hang; if that happens, re-run the job. Enabling Dynamo logs tends to reduce stalls: export TORCH_LOGS="+dynamo". To avoid repeated compilation cost across runs, enable the compile caches:
export TORCHINDUCTOR_FX_GRAPH_CACHE=1
export TORCHINDUCTOR_AUTOGRAD_CACHE=1Set these environment variables in your shell (or job script) before launching training.