PixelDiT

Training is managed via train.sh and follows a progressive three-stage pipeline to scale from 512px to 1024px resolution:

1. Stage 1 (Pre-training): 512×512 resolution from scratch. Uses fixed resolution and REPA loss for faster convergence.
2. Stage 2 (Multi-Aspect-Ratio Fine-tuning): 512×512 resolution. Enables multi-aspect ratio training and removes REPA loss (uses only flow matching loss).
3. Stage 3 (High-Res Fine-tuning): 1024×1024 resolution. Scales up to 1024px with multi-aspect ratio and no REPA loss.

To move between stages, use the --load_from flag to pass the checkpoint from the previous stage.

If you encounter sudden loss or gradient-norm spikes during training, you can enable post-modulation adaLN for the pixel-level (PiT) blocks.

By default, the model uses 6-way pre-modulation. Enabling post-modulation applies a 4-way scale/shift to the attention and MLP outputs (without a gate) for PiT blocks only. This is backward compatible and does not affect patch-level blocks.

Set pit_adaln_post_modulation: true in the denoiser config. Ready-to-use configs are available:

For 256×256:

bash train_c2i.sh --num-gpus 8 --config configs/pix256_xl_pit_post_modulation.yaml

For 512×512:

bash train_c2i.sh --num-gpus 8 --config configs/pix512_xl_pit_post_modulation.yaml

When configuring training, choose the correct flag based on your goal:

• --load_from: Use this to load weights only. This is used when transitioning between training stages (e.g., Stage 1 to Stage 2) where you want a fresh optimizer and scheduler.
• --resume_from: Use this to fully resume a training run. This restores the optimizer state, scheduler, and the current step count.

To resume training from a specific checkpoint, use the --ckpt-path flag. Note that auto_resume: true is enabled by default in the configs; if a checkpoint exists in the output directory, training will resume automatically. If you pass a filename to --ckpt-path, checkpoints will be auto-downloaded from HuggingFace if they are not found locally.

cd c2i/
bash train_c2i.sh --num-gpus 8 --config configs/pix256_xl.yaml \
  --ckpt-path /path/to/checkpoint.ckpt

Training is performed using the train_c2i.sh script. You can specify the number of GPUs and the configuration file.

To train on ImageNet 256×256:

cd c2i/
bash train_c2i.sh --num-gpus 8 --config configs/pix256_xl.yaml

To train on ImageNet 512×512:

cd c2i/
bash train_c2i.sh --num-gpus 8 --config configs/pix512_xl.yaml

To evaluate the model, use main.py predict to generate 50,000 images. The images and an output.npz file will be saved in c2i/train_logs/.

Example for ImageNet 256×256 (Epoch 80):

cd c2i/
torchrun --nproc_per_node=8 main.py predict \
  -c configs/pix256_xl.yaml \
  --ckpt_path=imagenet256_pixeldit_xl_epoch80.ckpt \
  --model.diffusion_sampler.class_path=src.diffusion.FlowDPMSolverSampler \
  --model.diffusion_sampler.init_args.num_steps=100 \
  --model.diffusion_sampler.init_args.guidance=3.25 \
  --model.diffusion_sampler.init_args.timeshift=1.0 \
  --model.diffusion_sampler.init_args.guidance_interval_min=0.1 \
  --model.diffusion_sampler.init_args.guidance_interval_max=1.0 \
  --per_run_seed=false --seed_everything=5000

To fully resume training (including optimizer state, scheduler, and step count), use the --resume_from flag. This is different from --load_from, which only loads model weights with a fresh optimizer.

bash train.sh configs/PixelDiT_1024px_pixel_diffusion_stage3.yaml \
  --resume_from=/path/to/checkpoint.pth \
  --work_dir=/path/to/output \
  --name=pixeldit-t2i-1024-stage3

To set up the PixelDiT environment, it is recommended to use the following Docker image:

nvcr.io/nvidia/pytorch:24.09-py3

Once inside the container or your preferred environment, install the required dependencies using pip:

pip install -r requirements.txt

You can download pre-trained PixelDiT-XL models for class-conditioned image generation on ImageNet from HuggingFace. Available resolutions include 256×256 and 512×512.

| Model | Epochs | Resolution | gFID | Checkpoint |
|:---:|:---:|:---:|:---:|:---:|
| PixelDiT-XL | 80  | 256×256 | 2.36 | [🤗 HuggingFace](https://huggingface.co/nvidia/PixelDiT-ImageNet/resolve/main/imagenet256_pixeldit_xl_epoch80.ckpt) |
| PixelDiT-XL | 160 | 256×256 | 1.97  | [🤗 HuggingFace](https://huggingface.co/nvidia/PixelDiT-ImageNet/resolve/main/imagenet256_pixeldit_xl_epoch160.ckpt) |
| PixelDiT-XL | 320 | 256×256 | 1.61 | [🤗 HuggingFace](https://huggingface.co/nvidia/PixelDiT-ImageNet/resolve/main/imagenet256_pixeldit_xl_epoch320.ckpt) |
| PixelDiT-XL | 850 | 512×512 | 1.81 | [🤗 HuggingFace](https://huggingface.co/nvidia/PixelDiT-ImageNet/resolve/main/imagenet512_pixeldit_xl.ckpt) |

Data preparation requirements differ by resolution:

• ImageNet 256×256: Follow the instructions in the REPA-E repository. Download and extract the ImageNet-1K training split, then run the preprocessing script provided there.
• ImageNet 512×512: Use the dataset_tool.py from the EDM2 repository to prepare the data.

Post-processing: After preprocessing, you must update the data_dir field in your configuration YAML (e.g., configs/pix256_xl.yaml or configs/pix512_xl.yaml) to point to your processed data directory.