Summary: We propose Explicit ShortCut (ESC), a framework that provides theoretical justification for the validity of shortcut models and disentangles concrete component-level choices, thereby enabling systematic identification of improvements. With our proposed improvements, the resulting one-step model achieves a new state-of-the-art FID50k of 2.85 on ImageNet-256×256, and further reaches FID50k of 2.53 with 2× training steps, under the classifier-free guidance setting without pre-training, distillation, or curriculum learning.

This implementation utilizes LMDB datasets with VAE-encoded latent representations for efficient training. The preprocessing pipeline is a reimplementation of the MAR. Once the ImageNet is downloaded in "YOUR/IMAGNET/PATH", Run the following to create the LMDB datasets:

torchrun preprocess_scripts/main_cache_imagenet.py \
--folder_dir "YOUR/IMAGNET/PATH/train"
--target_lmdb "YOUR/DESTINATION/LMDB/PATH"

See ./scripts for detailed training commands.

We provide pretrained checkpoints for models trained with class-consistent minibatching:

See ./scripts/run_baseline.sh

For the trained checkpoints, or the downloaded ones (.pt file), we provide a distributed evaluation script for large-scale sampling and quantitative evaluation (FID, IS):

torchrun --nproc_per_node=8 --nnodes=1 evaluate.py \
    --ckpt "/PATH/TO/THE/CHECKPOINTS" \
    --model "SiT-B/2" \
    --resolution 256 \
    --cfg-scale 1.0 \
    --per-proc-batch-size 128 \
    --num-fid-samples 50000 \
    --sample-dir "./fid_dir" \
    --compute-metrics \
    --num-steps 1 \
    --fid-statistics-file "./fid_stats/adm_in256_stats.npz" \
    --adapt-model

If there is any data type problem, it means that the numpy or torch version is not correct, you can run the following instead:

torchrun --nnodes=1 evaluate.py \
    --ckpt "/PATH/TO/THE/CHECKPOINTS" \
    --model "SiT-B/2" \
    --resolution 256 \
    --cfg-scale 1.0 \
    --per-proc-batch-size 128 \
    --num-fid-samples 50000 \
    --sample-dir "./fid_dir" \
    --compute-metrics \
    --num-steps 1 \
    --fid-statistics-file "./fid_stats/adm_float32_in256_stats.npz" \
    --adapt-model

This codebase is built upon REPA. We thank the authors for their excellent work and open-source contribution.

We also thank the original MeanFlow implementation: Gsunshine/MeanFlow, Gsunshine/py-meanflow, and zhuyu-cs/MeanFlow for their PyTorch reimplementation, which helped with early code restructuring.

For IMM, sCT, and CM, we thanks their (re-)implementation for our further remodularizing.

If you find our work is helpful to your research, please cite the following:

@inproceedings{lin2025designonestepdiffusionshortcutting,
      title={On the Design of One-step Diffusion via Shortcutting Flow Paths}, 
      author={Haitao Lin and Peiyan Hu and Minsi Ren and Zhifeng Gao and Zhi-Ming Ma and Guolin Ke and Tailin Wu and Stan Z. Li},
      booktitle={The Fourteenth International Conference on Learning Representations},
      year={2026},
      url={https://openreview.net/forum?id=k6q8rRYVQR}
}