Project Website | ArXiv Preprint

Official code implementation for "Distilling Diversity and Control in Diffusion Models".

Distilled diffusion models generate images in far fewer timesteps but suffer from "mode collapse" - producing similar outputs despite different random seeds. Our work addresses this critical limitation through:

1. Control Distillation: We discover that control mechanisms (Concept Sliders, LoRAs, DreamBooth) trained on base models can be directly applied to distilled models without retraining.

2. DT-Visualization: A novel analysis technique that reveals what diffusion models "think" the final image will be at intermediate steps.

3. Diversity Distillation: A hybrid inference approach using the base model for only the first timestep before switching to the distilled model, restoring diversity while maintaining speed.

conda create -n distillation python=3.9
conda activate distillation

git clone https://github.com/rohitgandikota/distillation.git
cd distillation
pip install -r requirements.txt

@article{gandikota2025distilling,
  title={Distilling Diversity and Control in Diffusion Models},
  author={Rohit Gandikota and David Bau},
  journal={arXiv preprint arXiv:2503.10637}
  year={2025}
}