| --- |
| 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} |
| } |
| ``` |