---

🛠️ Preparation

1. Installation

Set up the environment and install dependencies:

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

# Create and activate conda environment
conda create -n rectok python=3.10 -y
conda activate rectok

# Install requirements
pip install -r requirements.txt

2. Download Models

Download pretrained models and necessary data assets:

# Download from HuggingFace
huggingface-cli download QingyuShi/RecTok --local-dir ./pretrained_models
# Organize data assets and offline models
mv ./pretrained_models/data ./data
mv ./pretrained_models/offline_models.zip ./offline_models.zip
unzip offline_models.zip && rm offline_models.zip

3. Download ImageNet-1K

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

data/
├── fid_stats/                          # FID statistics files
│   ├── 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
└── val.txt                             # Validation file list

---

📊 Evaluation

Tokenizer Evaluation

Evaluate the reconstruction performance of the tokenizer:

bash run_eval_tokenizer.sh pretrained_models/RecTok_decft.pth   # path to RecTok checkpoint

Generative Model Evaluation

Evaluate the generation quality (FID, etc.), you can find the evaluation results in dir ./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 256x256:

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

---

🚀 Training

1. Tokenizer Training

Stage 1: Train RecTok Tokenizer

Please modify the --entity "YOUR_WANDB_ENTITY" if you want to use wandb. Otherwise please remenber to remove the --enable_wandb.

bash run_train_tokenizer.sh

Stage 2: Decoder Fine-tuning Run the following command for decoder fine-tuning:

bash run_decoder_finetune_tokenizer.sh <exp_name in Stage 1 run_train_tokenizer.sh>

2. Generative Model Training

Option A: Train from Scratch: Train the diffusion transformer model ($\text{DiT}^{\text{DH}}\text{-XL}$):

bash run_train_diffusion.sh <exp_name in Stage 2 run_decoder_finetune_tokenizer.sh>

Option B: Train with Pretrained RecTok: To train DiT based on our official pretrained RecTok 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{rectok,
  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={arXiv preprint arXiv:2512.13421},
  year={2025}
}

---

🙏 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.