BudgetPix-pMF

🌐 Project page  |  💻 Code  |  📄 Paper  |  🧩 BudgetAny

One checkpoint, any token budget. One-step class-conditional ImageNet generation in pixel space with pixel MeanFlow (pMF), fine-tuned with BudgetPix compute-adaptive tokenization. At inference you choose how many tokens the transformer runs on: an entropy-based detail map merges the flattest regions into coarse tokens first and gives the fine tokens to the most detailed ones; the layout is chosen up front, by default from the model's own dense sample of the same class and noise. The same weights serve every budget.

Unofficial, open-source implementation, taken from budgetpix/BudgetPix-public: a public reimplementation of BudgetPix: compute-adaptive tokenization for pixel-space image diffusion (JiT, MiniT2I, PixelDiT, pMF). This model is part of BudgetAny, a community reimplementation of BudgetPix (Kara et al., 2026). It is not an official release. Thanks to the BudgetPix authors for the method.

Files

file	base model	dense tokens	parameters	size	SHA-256
budgetpix_pmf_l16_256.pt	pMF-L/16 (one-step), 256 px	256	517 M	2.07 GB	dec78c17f667eefa…

SHA256SUMS lists the full checksums. Also in the repository: pmf_released_L16_full_uv.pt and pmf_released_B16_full_uv.pt, the released pMF weights converted to torch with the auxiliary v-branch.

Usage

git clone https://github.com/BudgetAny/BudgetPix && cd BudgetPix/pmf    # setup: see README.md
hf download BudgetAny/BudgetPix-pMF budgetpix_pmf_l16_256.pt --local-dir weights
python sample.py --ckpt weights/budgetpix_pmf_l16_256.pt --budget 128 --classes 207,88,417,933

--budget is the number of tokens per image (256 is the dense grid).

Citation

If you use this model or the code, please cite the BudgetPix paper:

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

Built on pixel MeanFlow and APT.

License

MIT, like the pMF weights it derives from.