BudgetPix: Compute-Adaptive Tokenization for Pixel-Space Image Diffusion

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, 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

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), which permits non-commercial research use only.

Citation

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