vita-epfl/rdm

RDM, short for Representation Distribution Matching, is the research code behind a paper on training AI image generators that can produce a finished image in a single step, instead of the many refinement steps most image generators use today. The authors are from EPFL, Valeo.ai, and Sorbonne University, and the project includes a research paper, a live demo, and downloadable model checkpoints. The core idea is to train the one step generator by matching the statistical distribution of features from generated images against features from real images, using a collection of frozen, already trained image encoders as the comparison tool. There is no separate teacher model guiding the training, no adversarial network fighting against the generator, and no need to unroll multiple sampling steps during training. The README reports that this method reaches state of the art results on a standard one step ImageNet benchmark, and can also convert an existing four step image model called FLUX.2 klein into a one step version that scores higher than the original on a standard text to image evaluation. Using this project requires installing it with pip, along with libraries such as PyTorch, timm, and Hugging Face's transformers and diffusers packages, and you must already have the relevant image datasets, like ImageNet or COCO, downloaded on your own disk since the code does not fetch them for you. The repository includes scripts to download pretrained checkpoints, run evaluation, reproduce the paper's figures, and train new models across multiple GPUs. This is a research codebase aimed at people working on generative image models, not a general purpose tool, and it requires meaningful GPU compute to train from scratch. It is released under the MIT license.