# 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]](https://arxiv.org/abs/2410.20202)


## 😀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.

![image](diagram.png)

## 🍉Requirement

```shell
pip install -r requirements.txt
```

## 🐬Preparation

### Clone

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

### Create an anaconda environment [Optional]:

```shell
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](https://cocodataset.org/).
* Extract them to the `data` folder.
* The directory structure will be as follows:

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

### Usage

#### Training

```shell
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:

```latex
@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},
}
```