Skip to content
/ SPM Public
forked from Con6924/SPM

Official implementation of paper "One-dimensional Adapter to Rule Them All: Concepts, Diffusion Models and Erasing Applications".

License

Notifications You must be signed in to change notification settings

civitai/SPM

 
 

Repository files navigation

CivitAI Semi-Permeable Membrane (SPM) Training

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.

Our Approach with SPMs

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.

Creating Your Composite SPM

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.

Steps to Create a Composite SPM:

  1. Train Individual Models: Start by training separate SPMs on different concept. Each model specializes in recognizing and moderating specific content types.
  2. Merge Models: Combine these trained models using a methodology akin to LoRA merging.
  3. 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.

Training Recommended VRAM

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.

Authors README

Concept Semi-Permeable Membrane

✏️ Project Page | 📄 arXiv | 🤗 Hugging Face

sample

(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.

Getting Started

0. Setup

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.

1. Training SPMs

In the demo.ipynb notebook we provide tutorials for setting up configs and training. Please refer to the notebook for further details.

2. Generate with SPMs

We provide three approaches to generate images after acquiring SPMs:

(Recommended) Generate with our provided code

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

Generate in the notebook demo

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.

Generate with SD-WebUI

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.

3. Evaluate SPMs

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}

Model Zoo

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.

References

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:

In addition, these repos inspires the implementation of ours:

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}
}

About

Official implementation of paper "One-dimensional Adapter to Rule Them All: Concepts, Diffusion Models and Erasing Applications".

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Python 82.2%
  • Jupyter Notebook 17.8%