Using the original code provided by the authors of the One-Dimensional Adapter to Rule Them All: Concepts, Diffusion Models and Erasing Applications paper, this repository equips you with configuration files to train SPMs on various diffusion models like SD1.5, SDXL, and Pony Diffusion.
At CivitAI, we leverage SPMs for content moderation to ensure that no CSAM or Toxic Mature Content is generated via our on-site generator. Our unique approach involves training multiple SPMs on distinct concepts and then merging these models into one just like you would do with multiple LoRAs. This method enhances the adaptability and effectiveness of our models in handling diverse content moderation needs.
To emulate CivitAI's SPMs setup, you should train individual SPMs on various concepts initially. Post-training, these models can be merged to form a comprehensive, unified model capable of sophisticated content moderation across different scenarios and diffusion models.
- Train Individual Models: Start by training separate SPMs on different concept. Each model specializes in recognizing and moderating specific content types.
- Merge Models: Combine these trained models using a methodology akin to LoRA merging.
- Evaluate and Iterate: Test the combined model's effectiveness across scenarios that the individual models were trained on, and iteratively refine the merging process for improved performance.
By following these steps, you can create a robust and versatile SPM that mirrors the functionality of CivitAI’s own systems, ensuring effective moderation across multiple diffusion platforms.
Base Model | Recommended VRAM |
---|---|
SD 1.5 | 20GB |
SDXL | 48GB |
We hope this guide assists you in your endeavors to develop advanced content moderation tools. For further assistance or collaboration inquiries, please feel free to contact us.
✏️ 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}
}