"ERASING CONCEPTS FROM DIFFUSION MODELS"
project: https://erasing.baulab.info
code: https://github.com/rohitgandikota/erasing
arxiv: https://arxiv.org/pdf/2303.07345.pdf

The folder `NSFW` contains the UNET network weights for fine-tuned stable-diffusion-v1.4 with nudity erased
	- diffusers-nudity-ESDu1-UNET.pt

The folder `art` contains the UNET network weights for fine-tuned stable-diffusion-v1.4 with several art styles erased
	- diffusers-VanGogh-ESDx1-UNET.pt
	- diffusers-KellyMcKernan-ESDx1-UNET.pt
	- diffusers-TylerEdlin-ESDx1-UNET.pt
	- diffusers-KilianEng-ESDx1-UNET.pt
	- diffusers-ThomasKinkade-ESDx1-UNET.pt
	- diffusers-AjinDemiHuman-ESDx1-UNET.pt

All the files are Huggingface's diffuser compatible: 

For complete code to generate images using these models please refer to our repository: https://github.com/rohitgandikota/erasing

Depending on the file-directory and use-case, feel free to edit the code available in https://github.com/rohitgandikota/erasing/blob/main/eval-scripts/generate-images.py

Quick-Start:
 
To load the weights use the following code 
```
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler
import torch

dir_ = "CompVis/stable-diffusion-v1-4"
vae = AutoencoderKL.from_pretrained(dir_, subfolder="vae")
tokenizer = CLIPTokenizer.from_pretrained(dir_, subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained(dir_, subfolder="text_encoder")
unet = UNet2DConditionModel.from_pretrained(dir_, subfolder="unet")

model_path = f'diffusers-VanGogh-ESDx1-UNET.pt'
unet.load_state_dict(torch.load(model_path))

```