Latent Watermark

Latent Watermark: Inject and Detect Watermarks in Latent Diffusion Space

Authors: Zheling Meng, Bo Peng, Jing Dong

New Laboratory of Pattern Recognition (NLPR), Institute of Automation, CAS

Setup

Install packages

conda create -n LWenv python=3.7.16
conda activate LWenv
pip install -r ./requirements.txt

Train the modules

0. Prepare for training

a. Training data Please prepare the training data before start. You can download LAION-Aesthetics-5+ same as our paper, or use your own data. And write your data path in the key "data_json" in the config file.

b. Pretrained SD Please download the pretrained v1.4 model into the folder "sd_ckpts" from Stable Diffusion. If you want to use other versions, modify the related configs in the config file as well.

1. Training stage 1

sh ./scripts/training_stage1.sh $cuda_device

2. Training stage 2

sh ./scripts/training_stage2.sh $cuda_device

3. Training stage 3

sh ./scripts/training_stage3.sh $cuda_device

Generate images

1. Inject watermarks

sh ./scripts/inject.sh $cuda_device

2. Extract watermarks

sh ./scripts/extract.sh $cuda_device

Evaluate the performance

Please refer to the repo WatermarkAttacker.

Acknowledgments

The code is built upon Stable Diffusion.

To-Do-List

• Fix the training and extraction bugs.
• Align the Performance.

Reference

Please cite our paper if you use our models in your works:

@article{meng2024latent,
  title={Latent Watermark: Inject and Detect Watermarks in Latent Diffusion Space},
  author={Meng, Zheling and Peng, Bo and Dong, Jing},
  journal={arXiv preprint arXiv:2404.00230},
  year={2024}
}