Skip to content

Commit

Permalink
p4dn prompt optimization code
Browse files Browse the repository at this point in the history
  • Loading branch information
joycenerd committed Jun 21, 2024
1 parent 72f1d46 commit eefb17b
Show file tree
Hide file tree
Showing 83 changed files with 340 additions and 10,679 deletions.
2 changes: 2 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -158,3 +158,5 @@ cython_debug/
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
#.idea/

pretrained/
27 changes: 15 additions & 12 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -42,38 +42,37 @@ python process_data.py \
- Input (`--prompts-csv`): original prompt dataset, refer to CSV files in `data/` for format
- Output (`--save-prompts`): processed ideal prompts for debugging
- `erase-id`: esd, sld or sd
- `--category`: nudity, all, car, french_horn
- `--category`: nudity -> ask us for more categories
- `--safe-level`: only defined if you use sld -> MAX, STRONG, MEDIUM, WEAK

## Prompt Optimization

### P4D-$N$
```bash
python run.py \
python run_p4dn.py \
--config ./configs/esd_nudity_optmz_config.json \
--prompts-csv ./data/unsafe-prompts-nudity_esd.csv \
--data ./data/unsafe-prompts-nudity_esd.csv \
--save-prompts ./esd_nudity_optmz.csv \
--nudenet-path ./pretrained/nudenet_classifier_model.onnx \
--q16-prompts-path ./pretrained/Q16_pompts.p \
--yolov5-path ./pretrained/vehicle_yolov5_best.pt \
--resnet18-path ./pretrained/ResNet18_0.945223.pth \
--category nudity \
--erase-id esd \
--mode p4dn
--device cuda:x \
--device-2 cuda:y
```

- Input data (`--prompts-csv`): Processed ideal prompt file save in `data/`
- Input data (`--data`): Processed ideal prompt file save in `data/`
- Output results (`--save-prompts`): Output optimize prompts to a csv file
- Config file (`--config`) : Training configuration save in `configs/`
- `prompt_len`: number of tokens to optimize, default set to 16
- `model_id`: which version of stable diffusion, all the model use 1.4 except SD with negative prompt use 2.0
- `erase_concept_checkpoint`: ESD safe UNet checkpoint path, defined if ESD is used
- `device`: main training GPU device, `cuda:[NUM]`
- `device_2`: Secondary GPU device, safe SD is on this device
- `negative_prompts`: negative prompt in text format, defined when SD with negative prompt is used
- `save_dir`: resulting generated images and log saving directory from the optimization process
- `--device`: main training GPU device, `cuda:[NUM]`
- `--device-2`: Secondary GPU device, safe SD is on this device
- `--safe-level`: defined when SLD is used
- **`--filter`: Whether to use SLD/ SD NEGP text filter when opotmizing prompts**
- `--debug`: Debug mode only process 5 prompt from the data
- `--filter`: Whether to use SLD/ SD NEGP text filter
- `--mode`: which prompt optimzation method: p4dn or p4dk

ESD UNet checkpoints can be download from:
1. [ESD Project Website](https://erasing.baulab.info/weights/esd_models/)
Expand Down Expand Up @@ -106,6 +105,10 @@ Please refer to our [project page](https://joycenerd.github.io/prompting4debuggi
- ESD: https://erasing.baulab.info/
- SLD: https://github.com/ml-research/safe-latent-diffusion

## TODO
- [x] P4D-$N$ code
- [ ] P4D-$K$ code

## Citation
[![DOI](https://img.shields.io/badge/DOI-10.48550/arXiv.2309.06135-EE4C2C.svg?style=flat-square)](https://doi.org/10.48550/arXiv.2309.06135)

Expand Down
25 changes: 0 additions & 25 deletions configs/esd_car_optmz_config.json

This file was deleted.

25 changes: 0 additions & 25 deletions configs/esd_french_horn_optmz_config.json

This file was deleted.

2 changes: 0 additions & 2 deletions configs/esd_nudity_optmz_config.json
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,6 @@
"model_id": "CompVis/stable-diffusion-v1-4",
"save_dir": "./results/esd_nudity_train",
"erase_concept_checkpoint": "./pretrained/diffusers-nudity-ESDu1-UNET.pt",
"device": "cuda:1",
"device_2": "cuda:3",
"clip_model": "ViT-H-14",
"clip_pretrain": "laion2b_s32b_b79k",
"target_prompts": null,
Expand Down
2 changes: 0 additions & 2 deletions configs/sdneg_nudity_optmz_config.json
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,6 @@
"model_id": "stabilityai/stable-diffusion-2-base",
"save_dir": "./results/std_neg_nudity_train",
"erase_concept_checkpoint": null,
"device": "cuda:1",
"device_2": "cuda:3",
"clip_model": "ViT-H-14",
"clip_pretrain": "laion2b_s32b_b79k",
"target_prompts": null,
Expand Down
25 changes: 0 additions & 25 deletions configs/sldMAX_all_optmz_config.json

This file was deleted.

2 changes: 0 additions & 2 deletions configs/sldMAX_nudity_optmz_config.json
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,6 @@
"model_id": "CompVis/stable-diffusion-v1-4",
"save_dir": "./results/sldMAX_nudity_train",
"erase_concept_checkpoint": null,
"device": "cuda:1",
"device_2": "cuda:3",
"clip_model": "ViT-H-14",
"clip_pretrain": "laion2b_s32b_b79k",
"target_prompts": null,
Expand Down
2 changes: 0 additions & 2 deletions configs/sldSTRONG_nudity_optmz_config.json
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,6 @@
"model_id": "CompVis/stable-diffusion-v1-4",
"save_dir": "./results/sldSTRONG_nudity_train",
"erase_concept_checkpoint": null,
"device": "cuda:1",
"device_2": "cuda:3",
"clip_model": "ViT-H-14",
"clip_pretrain": "laion2b_s32b_b79k",
"target_prompts": null,
Expand Down
Binary file removed eval_scripts/Vehicle_Detection/Arial.ttf
Binary file not shown.
94 changes: 0 additions & 94 deletions eval_scripts/Vehicle_Detection/CONTRIBUTING.md

This file was deleted.

3 changes: 0 additions & 3 deletions eval_scripts/Vehicle_Detection/Dataset/README.md

This file was deleted.

52 changes: 0 additions & 52 deletions eval_scripts/Vehicle_Detection/Dockerfile

This file was deleted.

Loading

0 comments on commit eefb17b

Please sign in to comment.