Compress a slow multi-step image generation model into a fast four-step version.
Fine-tune a distilled image generator using reward feedback to improve output quality.
Benchmark a few-step image generator against standard reward and quality metrics.
Reproduce the paper's results on models like SD3.5 and FLUX.2.
| harahan/rtdmd | hao0321/video-autopilot-kit | significant-gravitas/gravitasml | |
|---|---|---|---|
| Stars | 37 | 37 | 37 |
| Language | Python | Python | Python |
| Setup difficulty | hard | moderate | easy |
| Complexity | 5/5 | 3/5 | 2/5 |
| Audience | researcher | general | developer |
Figures from each repo's GitHub metadata at analysis time.
Requires GPU hardware and familiarity with training diffusion or flow-based image models, this is research code, not a packaged app.
RTDMD is a research code release accompanying an academic paper on training fast text-to-image generation models. Standard diffusion models generate images through many small steps, often 50 or more. This work focuses on training a model that produces good images in only four steps by combining distribution-matching distillation with reward-guided reinforcement learning. The training process works in two stages. The first stage, called AC-DMD, trains a student model to match the output distribution of a larger teacher model by computing a consistency loss across intermediate generation steps. The second stage, called RTDMD, continues training using a reinforcement learning approach that rewards the student model for producing outputs that score well on metrics like image quality, alignment with the text prompt, and human preference ratings. The two objectives are applied jointly so the student model improves on both dimensions at once. The results reported in the paper show the four-step RTDMD model surpassing its teacher on most evaluated metrics for three model families: SD3-M, SD3.5-M, and FLUX.2 4B. The FLUX.2 4B model trained with RTDMD scores higher than the larger FLUX.2 9B teacher on seven of nine metrics when both are compared at their standard settings, including improvements in image reward, aesthetic score, PickScore, and GenEval. The repository provides training scripts organized into two command-line trainers, one for each stage, along with YAML configuration files for five model backbones. Inference and reward evaluation scripts are also included. Pretrained checkpoints are published on Hugging Face. The code requires Python 3.10 or later and PyTorch.
Research code that compresses slow, many-step AI image generators into fast four-step versions while using reward feedback to keep or improve image quality.
Mainly Python. The stack also includes Python, PyTorch.
Setup difficulty is rated hard, with roughly 1day+ to a first successful run.
Mainly researcher.
This repo across BitVibe Labs
Verify against the repo before relying on details.