โ๏ธ Project Page | ๐ arXiv | ๐ค Hugging Face
(The generation samples demonstrating 'Cat' SPM, which is trained on SD v1.4 and applied on RealisticVision. The upper row shows the original generation and the lower row shows the SPM-equipped ones.)
We propose Concept Semi-Permeable Membrane (SPM), as a solution to erase or edit concepts for diffusion models (DMs).
Briefly, it can achieve two main purposes:
- Prevent the generation of target concept from the DMs, while
- Preserve the generation of non-target concept of the DMs.
SPM has following advantages:
- Data-free: no extra text or image data is needed for training SPM.
- Lightweight: the trainable parameters of the SPM is only 0.5% of the DM. A SPM for SD v1.4 only takes 1.7MB space for storage.
- Customizable: once obtained, SPMs of different target concept can be equipped simultaneously on the DM according to your needs.
- Model-transferable: SPM trained on certain DM can be directly transfered to other personalized models without additional tuning. e.g. SD v1.4 -> SD v1.5 / Dreamshaper-8 or any other similar community models.
We use Conda to setup the training environments:
conda create -n spm python=3.10
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
pip install xformers
pip install -r requirements.txt
Additionally, you can setup SD-WebUI for generation with SPMs as well.
In the demo.ipynb notebook we provide tutorials for setting up configs and training. Please refer to the notebook for further details.
We provide three approaches to generate images after acquiring SPMs:
First, you need to setup your generation config. configs/generaion.yaml is an example. Then you can generate images by running the following commands:
python infer_spm.py \
--config ${generation_config} \
--spm_path ${spm_1} ... ${spm_n} \
--base_model ${base_model_path_or_link} \ # e.g. CompVis/stable-diffusion-v1-4
The demo.ipynb notebook also offers codes for generate samples with single or multi SPMs.
Notice: In this way, the Facilitate Transfer mechanism of SPM will NOT be activated. SPMs will have a relatively higher impact on non-targeted concepts.
The SPMs can be well adapted to the SD-WebUI for more generation options. You can load SPM as a LoRA module to the desired SD model.
Notice: In this way, the Facilitate Transfer mechanism of SPM will NOT be activated. SPMs will have a relatively higher impact on non-targeted concepts.
To validate the provided results in our paper, you can run the following code to evaluate the trained SPMs on the four pre-defined tasks. To check the detailed arguments explaination, just run python evaluate_task.py -h
.
accelerate launch --num_processes ${num_gpus} evaluate_task.py \
--task ${task} \
--task_args ${task_args} \
--img_save_path ${img_save_path} \
--save_path ${save_path}
Trained SPM for SD v1.x:
Task Type | SPM |
---|---|
General Concepts | Snoopy, Mickey, Spongebob, Pikachu, Donald Duck, Cat, Wonder Woman (->Gal Gadot), Luke Skywalker (->Darth Vader), Joker (->Heath Ledger), Joker (->Batman) |
Artistic Styles | Van Gogh, Picasso, Rembrant, Comic |
Explicit Contents | Nudity |
SPM for SD v2.x and SDXL will be released in the future.
This repo is the code for the paper One-dimentional Adapter to Rule Them All: Concepts, Diffusion Models and Erasing Applications.
Thanks for the creative ideas of the pioneer researches:
- https://github.com/rohitgandikota/erasing: Erasing Concepts from Diffusion Models
- https://github.com/nupurkmr9/concept-ablation: Ablating Concepts in Text-to-Image Diffusion Models
- https://github.com/clear-nus/selective-amnesia: Selective Amnesia: A Continual Learning Approach for Forgetting in Deep Generative Models
In addition, these repos inspires the implementation of ours:
- https://github.com/p1atdev/LECO: Low-rank adaptation for Erasing COncepts from diffusion models
- https://github.com/cloneofsimo/lora: Low-rank Adaptation for Fast Text-to-Image Diffusion Fine-tuning
- https://github.com/kohya-ss/sd-scripts: Training, generation and utility scripts for Stable Diffusion
If you find this repo useful, you can cite our work as follows:
@misc{lyu2023onedimensional,
title={One-dimensional Adapter to Rule Them All: Concepts, Diffusion Models and Erasing Applications},
author={Mengyao Lyu and Yuhong Yang and Haiwen Hong and Hui Chen and Xuan Jin and Yuan He and Hui Xue and Jungong Han and Guiguang Ding},
year={2023},
eprint={2312.16145},
archivePrefix={arXiv},
primaryClass={cs.CV}
}