BudgetPix-JiT

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

One checkpoint, any token budget. Class-conditional ImageNet generation in pixel space with JiT (Just image Transformers), fine-tuned with BudgetPix compute-adaptive tokenization. At inference you choose how many tokens the diffusion 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 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_jit_b16_256.pth	JiT-B/16, 256 px	256	134 M	0.53 GB	30bfc892a288eb39…
budgetpix_jit_b32_512.pth	JiT-B/32, 512 px	256	136 M	0.54 GB	da192f0a436a5a75…
budgetpix_jit_l16_256.pth	JiT-L/16, 256 px	256	462 M	1.85 GB	90a24ddff0e5bb91…
budgetpix_jit_l32_512.pth	JiT-L/32, 512 px	256	465 M	1.86 GB	f49743558503d2b9…

SHA256SUMS lists the full checksums. Each file holds the EMA weights of the full adapted model (backbone plus the added modules) in fp32.

Usage

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

--budget is the number of tokens per image (256 is the dense grid). Default guidance 3.0; at 25 % of the grid about 6.0 works best for the 512-px models.

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 JiT and APT.

License

MIT, like the JiT weights it derives from.