An efficient watermarking method for latent diffusion models via low-rank adaptation

Code for our paper "An efficient watermarking method for latent diffusion models via low-rank adaptation".

You can download the paper via: [ArXiv]

😀Summary

A lightweight parameter fine-tuning strategy with low-rank adaptation and dynamic loss weight adjustment enables efficient watermark embedding in large-scale models while minimizing impact on image quality and maintaining robustness.

🍉Requirement

pip install -r requirements.txt

🐬Preparation

Clone

git clone https://github.com/MrDongdongLin/EW-LoRA

Create an anaconda environment [Optional]:

conda create -n ewlora python==3.8.18
conda activate ewlora
pip install -r requirements.txt

Prepare the training data:

Download the dataset files here.
Extract them to the data folder.
The directory structure will be as follows:

coco2017
└── train
   ├── img1.jpg
   ├── img2.jpg
   └── img3.jpg
└── test
   ├── img4.jpg
   ├── img5.jpg
   └── img6.jpg

Usage

Training

cd ./watermarker/stable_signature
CUDA_VISIBLE_DEVICES=0 python train_SS.py --num_keys 1 \
--train_dir ./Datasets/coco2017/train2017 \
--val_dir ./Datasets/coco2017/val2017 \
--ldm_config ./watermarker/stable_signature/configs/stable-diffusion/v1-inference.yaml \
--ldm_ckpt ../models/ldm_ckpts/sd-v1-4-full-ema.ckpt \
--msg_decoder_path ../models/wm_encdec/hidden/ckpts/dec_48b_whit.torchscript.pt \
--output_dir ./watermarker/stable_signature/outputs/ \
--task_name train_SS_fix_weights \
--do_validation \
--val_frep 50 \
--batch_size 4 \
--lambda_i 1.0 --lambda_w 0.2 \
--steps 20000 --val_size 100 \
--warmup_steps 20 \
--save_img_freq 100 \
--log_freq 1 --debug

Citation

If this work is helpful, please cite as:

@article{linEfficientWatermarkingMethod2024,
  title     = {An Efficient Watermarking Method for Latent Diffusion Models via Low-Rank Adaptation},
  author    = {Lin, Dongdong and Li, Yue and Tondi, Benedetta and Li, Bin and Barni, Mauro},
  year      = {2024},
  month     = oct,
  number    = {arXiv:2410.20202},
  eprint    = {2410.20202},
}