simillusion/dit-s4-cifar10

DiT-S/4 on CIFAR-10 (class-conditional)

A demo checkpoint from Simulusion, an open-source PyTorch library that puts world models behind one API. It implements the paper Scalable Diffusion Models with Transformers (Peebles & Xie, ICCV 2023).

A Diffusion Transformer (DiT-S/4, 33M parameters). In Simulusion, DiT is a building block for world models, not a world model itself: it is not action-conditioned.

Results

Trained class-conditionally on CIFAR-10 (32x32) for 40 epochs, about two hours on a 4 GB RTX 3050 laptop GPU. Samples above use classifier-free guidance at scale 4 and 1000 DDPM steps.

This is a short demo run sized for a laptop GPU, meant to show what the model does. It is not a reproduction of the paper's results. More clips and learning curves are in the gallery.

Usage

pip install simulusion huggingface_hub

from simulusion.models.diffusion.DiT import DiT

model = DiT.from_pretrained("simillusion/dit-s4-cifar10")

Checkpoints are plain state dicts and load with torch.load(..., weights_only=True). They include optimizer state, so you can continue training from them.

Training

Reproduce this checkpoint from a clone of the repository:

python demos/dit_demo.py train --epochs 40 --batch-size 64

Citation

If you use this model, please cite the original paper:

@inproceedings{peebles2023dit,
  title     = {Scalable Diffusion Models with Transformers},
  author    = {Peebles, William and Xie, Saining},
  booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision},
  year      = {2023},
  url       = {https://arxiv.org/abs/2212.09748},
}

and Simulusion:

@software{simulusion,
  title  = {Simulusion: a PyTorch library for world models},
  author = {Thakkar, Param},
  year   = {2026},
  url    = {https://github.com/ParamThakkar123/simulusion},
}