Erasing Concepts from Diffusion Models

Project Website | Arxiv Preprint | Fine-tuned Weights | Demo

Updated code 🚀 - Shared diffusers training core

Code Update 🚀

We are releasing a cleaner code for ESD with diffusers support. Compared to our old-code this version uses almost half the GPU memory and is 5-8 times faster. The training entrypoints for SD, SDXL, and FLUX now share the same core logic, save metadata-aware checkpoints, and can target any compatible base model in the corresponding diffusers family via --basemodel_id.

To use the older version please visit our older commit

Installation Guide

We recently updated our codebase to be much more cleaner and faster. The setup is also simple

git clone https://github.com/rohitgandikota/erasing.git
cd erasing
conda create --name erasing python=3.14
conda activate erasing
pip install -r requirements.txt

for FLUX training:

pip install --upgrade diffusers transformers torch

Training Guide

SDv1.4

After installation, follow these instructions to train a custom ESD model for Stable Diffusion V1.x. Canonical train methods are 'esd-x', 'esd-u', 'selfattn', 'esd-all', and 'esd-x-strict'. Legacy aliases such as 'xattn', 'noxattn', and 'full' still work.

python esd_sd.py --erase_concept 'Van Gogh' --train_method 'esd-x'

💡 New application: You can now erase an attribute from a concept!! Instead of erasing a whole concept you can just precisely remove some of its attributes. For example, you can erase hats from cowboys but keep the rest intact!

python esd_sd.py --erase_concept 'cowboy hat' --erase_from 'cowboy' --train_method 'esd-x'

SDXL

After installation, follow these instructions to train a custom ESD model for SDXL. Pick from following 'esd-x', 'esd-x-strict', 'esd-u', 'esd-all' (NOTE: 'esd-u' and 'esd-all' are still experimental and might produce unexpected artifacts):

python esd_sdxl.py --erase_concept 'Van Gogh' --train_method 'esd-x-strict'

FLUX

After installation (make sure: diffusers and transformers are up-to-date), follow these instructions to train a custom ESD model for FLUX.1-dev. Pick from following 'esd-x', 'esd-x-strict':

python esd_flux.py --erase_concept 'monster' --train_method 'esd-x' --negative_guidance 2

FLUX.2 Klein

FLUX.2 Klein uses the newer Flux2KleinPipeline backbone, so make sure your diffusers and transformers installs are recent enough for FLUX.2 support. The default base model is black-forest-labs/FLUX.2-klein-base-4B, and the new entrypoint keeps the same ESD loss while following the Klein tokenizer / transformer / scheduler flow:

python esd_flux2_klein.py --erase_concept 'monster' --train_method 'esd-x' --negative_guidance 2

All four training scripts also accept --basemodel_id so you can point them at any compatible diffusers model in that family.

The optimization process for erasing undesired visual concepts from pre-trained diffusion model weights involves using a short text description of the concept as guidance. The ESD model is fine-tuned with the conditioned and unconditioned scores obtained from frozen SD model to guide the output away from the concept being erased. The model learns from it's own knowledge to steer the diffusion process away from the undesired concept.

Generating Images

Generating images from a custom ESD model is super easy. Please follow the notebooks in notebooks/.

For an automated script to generate a ton of images for your evaluations use our evalscripts. The loader now auto-detects whether a checkpoint targets unet or transformer, so the same script works for SD, SDXL, and FLUX checkpoints.

python evalscripts/generate-images.py --base_model 'stabilityai/stable-diffusion-xl-base-1.0' --esd_path 'esd-models/sdxl/esd-kelly-from-kelly.safetensors' --num_samples 1 --prompts_path 'data/kelly_prompts.csv' --num_inference_steps 20 --guidance_scale 7

UPDATE (NudeNet)

If you want to recreate the results from our paper on NSFW task - please use this https://drive.google.com/file/d/1J_O-yZMabmS9gBA2qsFSrmFoCl7tbFgK/view?usp=sharing

• Untar this file and save it in the homedirectory '~/.NudeNet'
• This should enable the right results as we use this checkpoint for our analysis.

Running Gradio Demo Locally

To run the gradio interactive demo locally, clone the files from demo repository

• Create an environment using the packages included in the requirements.txt file
• Run python app.py
• Open the application in browser at http://127.0.0.1:7860/
• Train, evaluate, and save models using our method

NOTE ON LICENSE

The code and methods behind our work have been released under MIT. However, the models that you use our methods with, might be on a different licenses. Please read the model's license (the model you are using) carefully for more details.

Citing our work

The preprint can be cited as follows

@inproceedings{gandikota2023erasing,
  title={Erasing Concepts from Diffusion Models},
  author={Rohit Gandikota and Joanna Materzy\'nska and Jaden Fiotto-Kaufman and David Bau},
  booktitle={Proceedings of the 2023 IEEE International Conference on Computer Vision},
  year={2023}
}