RecTok: Reconstruction Distillation along Rectified Flow [CVPR 2026]

Official PyTorch Implementation

---

⚡ Quick Start

Already have the environment set up? Download models and run evaluation in three commands:

# 1. Download pretrained models and data assets
huggingface-cli download QingyuShi/RecTok --local-dir ./pretrained_models
mv ./pretrained_models/data ./data
mv ./pretrained_models/offline_models.zip ./offline_models.zip && unzip offline_models.zip && rm offline_models.zip

# 2. Evaluate tokenizer (results saved to ./work_dirs/tokenizer_training/RecTok_eval/)
bash run_eval_tokenizer.sh pretrained_models/RecTok_decft.pth

# 3. Evaluate generative model (results saved to ./work_dirs/gen_model_training/RecTok_eval/)
bash run_eval_diffusion.sh \
    pretrained_models/RecTok_decft.pth \
    pretrained_models/ditdhxl_epoch_0599.pth \
    pretrained_models/ditdhs_epoch_0029.pth

---

🎨 Demo (No ImageNet Required)

Generate images from class labels with a single GPU — no dataset needed:

# Basic: 8 images across 8 diverse classes
python demo.py \
    --tokenizer_path pretrained_models/RecTok_decft.pth \
    --model_path pretrained_models/ditdhxl_epoch_0599.pth

# Recommended: enable auto-guidance for best quality (cfg=1.29)
python demo.py \
    --tokenizer_path pretrained_models/RecTok_decft.pth \
    --model_path pretrained_models/ditdhxl_epoch_0599.pth \
    --auto_guidance_path pretrained_models/ditdhs_epoch_0029.pth \
    --cfg 1.29

# Custom classes: 4 images each for golden retriever, macaw, and balloon
python demo.py \
    --tokenizer_path pretrained_models/RecTok_decft.pth \
    --model_path pretrained_models/ditdhxl_epoch_0599.pth \
    --classes 207 88 417 \
    --num_per_class 4

Latent stats: The first time you run demo.py, if work_dirs/stats.pkl is not present, generation still works but quality may be slightly lower. Run the full eval script once to create the cache, or copy stats.pkl from a completed run.

Output: Individual PNGs + grid.png are saved to ./demo_outputs/ by default.

---

🛠️ Preparation

1. Installation

# Clone the repository
git clone https://github.com/Shi-qingyu/RecTok.git
cd RecTok

# Create virtual environment and install all dependencies
uv sync

2. Download Models

Download pretrained models and necessary data assets (includes data/train.txt, data/val.txt, and FID statistics):

# Download from HuggingFace
huggingface-cli download QingyuShi/RecTok --local-dir ./pretrained_models

# Move data assets and offline models into place
mv ./pretrained_models/data ./data
mv ./pretrained_models/offline_models.zip ./offline_models.zip
unzip offline_models.zip && rm offline_models.zip

Model	Type	Params	Hugging Face
RecTok	Tokenizer	172M	🤗 rectok
RecTok-decft	Tokenizer	172M	🤗 rectok-decft
$\text{DiT}^{\text{DH}}\text{-XL}$-80e	Generator	839M	🤗 ditdh-xl-80e
$\text{DiT}^{\text{DH}}\text{-XL}$-600e	Generator	839M	🤗 ditdh-xl-600e
$\text{DiT}^{\text{DH}}\text{-S}$-30e (auto-guidance only)	Generator	193M	🤗 ditdh-s-30e

3. Download ImageNet-1K

Please download ImageNet-1K to ./data. Your directory structure should look like this:

data/
├── fid_stats/                          # FID statistics files (provided via HuggingFace download)
│   ├── adm_in256_stats.npz             # For gFID
│   ├── val_fid_statistics_file_256.npz # For rFID
│   └── val_fid_statistics_file_512.npz # For rFID
├── imagenet/                           # ImageNet dataset
│   ├── train/
│   │   ├── n01440764/
│   │   └── ...
│   └── val/
│       ├── n01440764/
│       └── ...
├── train.txt                           # Training file list (provided via HuggingFace download)
└── val.txt                             # Validation file list (provided via HuggingFace download)

---

📊 Evaluation

Tokenizer Evaluation

Evaluate the reconstruction performance (rFID, PSNR, SSIM, LPIPS). Results are saved to ./work_dirs/tokenizer_training/RecTok_eval/:

bash run_eval_tokenizer.sh pretrained_models/RecTok_decft.pth

Generative Model Evaluation

Evaluate the generation quality (gFID, Inception Score). Results are saved to ./work_dirs/gen_model_training/RecTok_eval/:

bash run_eval_diffusion.sh \
    pretrained_models/RecTok_decft.pth \        # path to RecTok checkpoint
    pretrained_models/ditdhxl_epoch_0599.pth \  # path to DiTDH-XL checkpoint
    pretrained_models/ditdhs_epoch_0029.pth     # path to autoguidance model checkpoint

Selected examples of class-conditional generation results on ImageNet-1K 256×256:

FID-50K and Inception Score without and with CFG:

cfg	Model	Epochs	FID-50K	Inception Score	#params
1.0	$\text{DiT}^{\text{DH}}\text{-XL}$ + RecTok	80	2.09	198.6	839M
1.29	$\text{DiT}^{\text{DH}}\text{-XL}$ + RecTok	80	1.48	223.8	839M
1.0	$\text{DiT}^{\text{DH}}\text{-XL}$ + RecTok	600	1.34	254.6	839M
1.29	$\text{DiT}^{\text{DH}}\text{-XL}$ + RecTok	600	1.13	289.2	839M

---

🚀 Training

WandB: All training scripts support --enable_wandb. Set --entity YOUR_WANDB_ENTITY in the script to log runs, or remove --enable_wandb to disable logging entirely.

exp_name: Each training stage creates a folder under work_dirs/ named by --exp_name (default: a timestamp like 20240501_1200_exp). This name is passed between stages to locate checkpoints.

1. Tokenizer Training

Stage 1: Train RecTok Tokenizer (~200 epochs, outputs to work_dirs/tokenizer_training/<exp_name>/):

bash run_train_tokenizer.sh

Stage 2: Decoder Fine-tuning (~100 epochs, pass the exp_name from Stage 1):

bash run_decoder_finetune_tokenizer.sh <exp_name from Stage 1>

2. Generative Model Training

Option A: Train from Scratch (uses the tokenizer from Stage 2):

bash run_train_diffusion.sh <exp_name from Stage 2>

Option B: Train with Pretrained RecTok (uses official pretrained weights):

mkdir -p work_dirs/tokenizer_training/rectok/checkpoints
cp pretrained_models/RecTok_decft.pth work_dirs/tokenizer_training/rectok/checkpoints/latest.pth
bash run_train_diffusion.sh rectok

---

📜 Citation

If you find this work useful for your research, please consider citing:

@article{rectokcvpr26,
  title={RecTok: Reconstruction Distillation along Rectified Flow},
  author={Shi, Qingyu and Wu, Size and Bai, Jinbin and Yu, Kaidong and Wang, Yujing and Tong, Yunhai and Li, Xiangtai and Li, Xuelong},
  journal={CVPR},
  year={2026}
}

---

🙏 Acknowledgements

We thank the authors of lDeTok, RAE, MAE, DiT, and LightningDiT for their foundational work.

Our codebase builds upon several excellent open-source projects, including lDeTok, RAE, and torch_fidelity. We are grateful to the communities behind them.

We sincerely thank Jiawei Yang and Boyang Zheng for providing insightful feedback.