Fractal Generative Models

This is a PyTorch/GPU implementation of the paper Fractal Generative Models:

@article{li2025fractal,
  title={Fractal Generative Models},
  author={Li, Tianhong and Sun, Qinyi and Fan, Lijie and He, Kaiming},
  journal={arXiv preprint arXiv:2502.17437},
  year={2025}
}

FractalGen enables pixel-by-pixel high-resolution image generation for the first time. This repo contains:

• 🪐 A simple PyTorch implementation of Fractal Generative Model.
• ⚡️ Pre-trained pixel-by-pixel generation models trained on ImageNet 64x64 and 256x256.
• 💥 A self-contained Colab notebook for running pre-trained models tasks.
• 🛸 A training and evaluation script using PyTorch DDP.

Preparation

Dataset

Download ImageNet dataset, and place it in your IMAGENET_PATH.

Installation

Download the code:

git clone https://github.com/LTH14/fractalgen.git
cd fractalgen

A suitable conda environment named fractalgen can be created and activated with:

conda env create -f environment.yaml
conda activate fractalgen

Download pre-trained models:

python util/download.py

For convenience, our pre-trained models can be downloaded directly here as well:

Usage

Demo

Run our interactive visualization demo using Colab notebook!

Training

The below training scripts have been tested on 4x8 H100 GPUs.

Example script for training FractalAR on ImageNet 64x64 for 800 epochs:

torchrun --nproc_per_node=8 --nnodes=4 --node_rank=${NODE_RANK} --master_addr=${MASTER_ADDR} --master_port=${MASTER_PORT} \
main_fractalgen.py \
--model fractalar_in64 --img_size 64 --num_conds 1 \
--batch_size 64 --eval_freq 40 --save_last_freq 10 \
--epochs 800 --warmup_epochs 40 \
--blr 5.0e-5 --weight_decay 0.05 --attn_dropout 0.1 --proj_dropout 0.1 --lr_schedule cosine \
--gen_bsz 256 --num_images 8000 --num_iter_list 64,16 --cfg 11.0 --cfg_schedule linear --temperature 1.03 \
--output_dir ${OUTPUT_DIR} --resume ${OUTPUT_DIR} \
--data_path ${IMAGENET_PATH} --grad_checkpointing --online_eval

Example script for training FractalMAR on ImageNet 64x64 for 800 epochs:

torchrun --nproc_per_node=8 --nnodes=4 --node_rank=${NODE_RANK} --master_addr=${MASTER_ADDR} --master_port=${MASTER_PORT} \
main_fractalgen.py \
--model fractalmar_in64 --img_size 64 --num_conds 5 \
--batch_size 64 --eval_freq 40 --save_last_freq 10 \
--epochs 800 --warmup_epochs 40 \
--blr 5.0e-5 --weight_decay 0.05 --attn_dropout 0.1 --proj_dropout 0.1 --lr_schedule cosine \
--gen_bsz 256 --num_images 8000 --num_iter_list 64,16 --cfg 6.5 --cfg_schedule linear --temperature 1.02 \
--output_dir ${OUTPUT_DIR} --resume ${OUTPUT_DIR} \
--data_path ${IMAGENET_PATH} --grad_checkpointing --online_eval

Example script for training FractalMAR-L on ImageNet 256x256 for 800 epochs:

torchrun --nproc_per_node=8 --nnodes=4 --node_rank=${NODE_RANK} --master_addr=${MASTER_ADDR} --master_port=${MASTER_PORT} \
main_fractalgen.py \
--model fractalmar_large_in256 --img_size 256 --num_conds 5 --guiding_pixel \
--batch_size 32 --eval_freq 40 --save_last_freq 10 \
--epochs 800 --warmup_epochs 40 \
--blr 5.0e-5 --weight_decay 0.05 --attn_dropout 0.1 --proj_dropout 0.1 --lr_schedule cosine \
--gen_bsz 256 --num_images 8000 --num_iter_list 64,16,16 --cfg 21.0 --cfg_schedule linear --temperature 1.1 \
--output_dir ${OUTPUT_DIR} --resume ${OUTPUT_DIR} \
--data_path ${IMAGENET_PATH} --grad_checkpointing --online_eval

Evaluation

Evaluate pre-trained FractalAR on ImageNet 64x64 unconditional likelihood estimation (single GPU):

torchrun --nproc_per_node=1 --nnodes=1 --node_rank=0 \
main_fractalgen.py \
--model fractalar_in64 --img_size 64 --num_conds 1 \
--nll_bsz 128 --nll_forward_number 1 \
--output_dir pretrained_models/fractalar_in64 \
--resume pretrained_models/fractalar_in64 \
--data_path ${IMAGENET_PATH} --seed 0 --evaluate_nll

Evaluate pre-trained FractalMAR on ImageNet 64x64 unconditional likelihood estimation (single GPU):

torchrun --nproc_per_node=1 --nnodes=1 --node_rank=0 \
main_fractalgen.py \
--model fractalmar_in64 --img_size 64 --num_conds 5 \
--nll_bsz 128 --nll_forward_number 10 \
--output_dir pretrained_models/fractalmar_in64 \
--resume pretrained_models/fractalmar_in64 \
--data_path ${IMAGENET_PATH} --seed 0 --evaluate_nll

Evaluate pre-trained FractalAR on ImageNet 64x64 class-conditional generation:

torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
main_fractalgen.py \
--model fractalar_in64 --img_size 64 --num_conds 1 \
--gen_bsz 512 --num_images 50000 \
--num_iter_list 64,16 --cfg 11.0 --cfg_schedule linear --temperature 1.03 \
--output_dir pretrained_models/fractalar_in64 \
--resume pretrained_models/fractalar_in64 \
--data_path ${IMAGENET_PATH} --seed 0 --evaluate_gen

Evaluate pre-trained FractalMAR on ImageNet 64x64 class-conditional generation:

torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
main_fractalgen.py \
--model fractalmar_in64 --img_size 64 --num_conds 5 \
--gen_bsz 1024 --num_images 50000 \
--num_iter_list 64,16 --cfg 6.5 --cfg_schedule linear --temperature 1.02 \
--output_dir pretrained_models/fractalmar_in64 \
--resume pretrained_models/fractalmar_in64 \
--data_path ${IMAGENET_PATH} --seed 0 --evaluate_gen

Evaluate pre-trained FractalMAR-Huge on ImageNet 256x256 class-conditional generation:

torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
main_fractalgen.py \
--model fractalmar_huge_in256 --img_size 256 --num_conds 5 --guiding_pixel \
--gen_bsz 1024 --num_images 50000 \
--num_iter_list 64,16,16 --cfg 19.0 --cfg_schedule linear --temperature 1.1 \
--output_dir pretrained_models/fractalmar_huge_in256 \
--resume pretrained_models/fractalmar_huge_in256 \
--data_path ${IMAGENET_PATH} --seed 0 --evaluate_gen

For ImageNet 256x256, the optimal classifier-free guidance values --cfg that achieve the best FID are 29.0 for FractalMAR-Base and 21.0 for FractalMAR-Large.

Acknowledgements

We thank Google TPU Research Cloud (TRC) for granting us access to TPUs, and Google Cloud Platform for supporting GPU resources.

Contact

If you have any questions, feel free to contact me through email (tianhong@mit.edu). Enjoy!