PRA / README.md
nielsr's picture
nielsr HF Staff
Add model card, link to paper and GitHub repository
27573c2 verified
|
Raw
History Blame
2.46 kB
---
pipeline_tag: unconditional-image-generation
---
# Parallel Rollout Approximation (PRA)
This repository contains the weights for **Parallel Rollout Approximation (PRA)**, a scalable framework for class-conditional pixel-space autoregressive image generation.
More details can be found in the paper [Parallel Rollout Approximation for Pixel-Space Autoregressive Image Generation](https://huggingface.co/papers/2606.27978).
* **Repository:** [GitHub Repository](https://github.com/MangataX/PRA)
* **Paper:** [arXiv:2606.27978](https://huggingface.co/papers/2606.27978)
## Model Description
Parallel Rollout Approximation (PRA) is a pixel-space continuous-token autoregressive (AR) generation model. PRA generates low-dimensional intermediate states instead of high-dimensional pixel patches, mapping them back to pixel-space tokens with a pixel decoder. It effectively mitigates error accumulation during autoregressive steps by approximating the pixel-feedback interface encountered during inference-time rollout while retaining parallel teacher-forced training.
## Model Checkpoints
The following checkpoints are available:
| Model | Params | FID (256x256) | Weight |
|:---:|:---:|:---:|:---:|
| PRA-S | 135M | 2.58 | [PRA_S.pt](https://huggingface.co/MangataX/PRA/blob/main/PRA_S.pt) |
| PRA-B | 250M | 2.21 | [PRA_B.pt](https://huggingface.co/MangataX/PRA/blob/main/PRA_B.pt) |
| PRA-L | 511M | 1.94 | [PRA_L.pt](https://huggingface.co/MangataX/PRA/blob/main/PRA_L.pt) |
## Environment & Usage
For environment setup, training, and evaluation scripts, please refer to the official [GitHub Repository](https://github.com/MangataX/PRA).
### Sampling Example
You can run distributed class-balanced sampling using the `sample_ddp.py` script provided in the repository:
```shell
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-png
```
## Citation
```bibtex
@article{xu2026parallel,
title={Parallel Rollout Approximation for Pixel-Space Autoregressive Image Generation},
author={Xu, Jiayi and He, Di and Ke, Guolin},
journal={arXiv preprint arXiv:2606.27978},
year={2026}
}
```