uses sdxl_vae from https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/tree/main

Paper link: https://arxiv.org/pdf/2407.15811

20 epochs of training on a 1b param model (common canvas cc-by dataset) using 8xH100 for 5 days, around $2000 of compute.

You can download weights here: https://huggingface.co/SwayStar123/MicroDiT/blob/main/no_cfg/microdit_model_epoch_19.pt and put them in models folder, and use test_model.ipynb to inference

50 epochs of training on a 150m param model (CelebA dataset)

Credit to original authors: Vikash Sehwag, Xianghao Kong, Jingtao Li, Michael Spranger, Lingjuan Lyu

Thank you to Baber from discord for saving me hours of my time