---
license: mit
library_name: pytorch
pipeline_tag: text-to-image
tags:
  - image-generation
  - diffusion
  - pixel-space
  - adaptive-tokenization
  - imagenet
  - text-to-image
  - one-step-generation
---

<h1 align="center">BudgetPix: Compute-Adaptive Tokenization for Pixel-Space Image Diffusion</h1>

<p align="center">
  <a href="https://arxiv.org/abs/2610.12307"><img src="https://img.shields.io/badge/arXiv-2610.12307-b31b1b.svg" alt="arXiv"></a>
  <a href="https://karaozgur.com/BudgetPix/"><img src="https://img.shields.io/badge/Project-Page-2ea44f.svg" alt="Project page"></a>
  <a href="https://karaozgur.com/BudgetPix/static/paper/BudgetPix.pdf"><img src="https://img.shields.io/badge/Paper-PDF-blue.svg" alt="Paper PDF"></a>
  <a href="https://github.com/budgetpix/BudgetPix-public"><img src="https://img.shields.io/badge/GitHub-Code-181717.svg?logo=github" alt="Code on GitHub"></a>
</p>

<p align="center"><img src="teaser.gif" width="720" alt="BudgetPix teaser: one text-to-image model run from 100 % down to 10 % of the tokens"></p>

Checkpoints of the public reimplementation of **BudgetPix**. Pixel-space diffusion transformers spend the same compute
on flat sky as on a face; BudgetPix lets one model run at many compute budgets. An entropy quadtree merges flat regions
of the image into coarser tokens, the backbone runs on the shorter sequence, and a small token-local refiner puts the
detail back inside each merged token. Every added module is zero-initialised, so the adapted model starts as the
released model; at inference a single checkpoint is driven with a token budget. The code that loads, samples from and
trains these files is at https://github.com/budgetpix/BudgetPix-public.

Every file is the inference-only EMA weights of the full adapted model (released backbone plus the added modules) in
fp32, with the geometry and trainer arguments it was trained with.

| file | model | task | resolution | dense tokens | size |
|---|---|---|---|---|---|
| `budgetpix_jit_b16_256.pth` | JiT-B/16 | class-conditional ImageNet | 256 | 256 | 0.53 GB |
| `budgetpix_jit_b32_512.pth` | JiT-B/32 | class-conditional ImageNet | 512 | 256 | 0.54 GB |
| `budgetpix_jit_l16_256.pth` | JiT-L/16 | class-conditional ImageNet | 256 | 256 | 1.85 GB |
| `budgetpix_jit_l32_512.pth` | JiT-L/32 | class-conditional ImageNet | 512 | 256 | 1.86 GB |
| `budgetpix_minit2i_b_512.pt` | MiniT2I-B/16 | text-to-image | 512 | 1024 | 1.04 GB |
| `budgetpix_minit2i_l_512.pt` | MiniT2I-L/16 | text-to-image | 512 | 1024 | 3.66 GB |
| `budgetpix_pixeldit_1024.pt` | PixelDiT-T2I 1.3B | text-to-image | 1024 | 4096 | 5.23 GB |
| `budgetpix_pmf_l16_256.pt` | pMF-L/16 | one-step class-conditional ImageNet | 256 | 256 | 2.07 GB |

Extras for the pMF line: `pmf_released_B16_full_uv.pt` and `pmf_released_L16_full_uv.pt` are the released pixel MeanFlow
weights ([Lyy0725/pMF](https://huggingface.co/Lyy0725/pMF), MIT) converted to PyTorch with the auxiliary v-branch that
fine-tuning needs; the Hub's inference files carry the u-path only. `SHA256SUMS` lists the hash of every file.

## Usage

```bash
git clone https://github.com/budgetpix/BudgetPix-public && cd BudgetPix-public
pip install -r requirements.txt
python scripts/download_checkpoints.py --all                 # fetches from this repo and checks the hashes

cd jit && python sample.py --ckpt ../checkpoints/budgetpix_jit_l32_512.pth --budget 128 --classes 207,88,417,933 --overlay
```

The MiniT2I, PixelDiT and pMF lines import the official PyTorch code of those models at run time (`MINIT2I_ROOT`,
`PIXELDIT_ROOT`, `PMF_ROOT`); the text encoders (`google/flan-t5-large`, `Efficient-Large-Model/gemma-2-2b-it`) are
downloaded from the Hub on first use. See the GitHub README for sampling, training and evaluation commands.

## License

The JiT, MiniT2I and pMF checkpoints derive from MIT-licensed weights and are released under MIT. The PixelDiT
checkpoint derives from NVIDIA's released PixelDiT weights and is subject to the
[NVIDIA License (NSCLv1)](https://github.com/NVlabs/PixelDiT/blob/master/LICENSE), which permits non-commercial
research use only.

## Citation

```bibtex
@article{kara2026budgetpix,
  title={BudgetPix: Compute-Adaptive Tokenization for Pixel-Space Image Diffusion},
  author={Ozgur Kara and Yujia Chen and Daniel Watson and David Forsyth and James Matthew Rehg and Wen-Sheng Chu and Du Tran},
  journal={arXiv preprint arXiv:2610.12307},
  year={2026}
}
```