Spaces:
Running
on
Zero
A newer version of the Gradio SDK is available:
5.9.1
title: Unboxing SDXL with SAEs
app_file: app.py
sdk: gradio
sdk_version: 4.44.1
Unpacking SDXL Turbo: Interpreting Text-to-Image Models with Sparse Autoencoders
This repository contains code to reproduce results from our paper (https://arxiv.org/abs/2410.22366) on using sparse autoencoders (SAEs) to analyze and interpret the internal representations of text-to-image diffusion models, specifically SDXL Turbo.
Repository Structure
|-- SAE/ # Core sparse autoencoder implementation
|-- SDLens/ # Tools for analyzing diffusion models
| `-- hooked_sd_pipeline.py # Modified stable diffusion pipeline
|-- scripts/
| |-- collect_latents_dataset.py # Generate training data
| `-- train_sae.py # Train SAE models
|-- utils/
| `-- hooks.py # Hook utility functions
|-- checkpoints/ # Pretrained SAE model checkpoints
|-- app.py # Demo application
|-- app.ipynb # Interactive notebook demo
|-- example.ipynb # Usage examples
`-- requirements.txt # Python dependencies
Installation
pip install -r requirements.txt
Demo Application
You can try our gradio demo application (app.ipynb
) to browse and experiment with 20K+ features of our trained SAEs out-of-the-box. You can find the same notebook on Google Colab.
Usage
- Collect latent data from SDXL Turbo:
python scripts/collect_latents_dataset.py --save_path={your_save_path}
Train sparse autoencoders:
2.1. Insert the path of stored latents and directory to store checkpoints in
SAE/config.json
2.2. Run the training script:
python scripts/train_sae.py
Pretrained Models
We provide pretrained SAE checkpoints for 4 key transformer blocks in SDXL Turbo's U-Net. See example.ipynb
for analysis examples and visualization of learned features.
Citation
If you find this code useful in your research, please cite our paper:
@misc{surkov2024unpackingsdxlturbointerpreting,
title={Unpacking SDXL Turbo: Interpreting Text-to-Image Models with Sparse Autoencoders},
author={Viacheslav Surkov and Chris Wendler and Mikhail Terekhov and Justin Deschenaux and Robert West and Caglar Gulcehre},
year={2024},
eprint={2410.22366},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2410.22366},
}
Acknowledgements
The SAE component was implemented based on openai/sparse_autoencoder
repository.