Awesome
AquaLoRA: Toward White-box Protection for Customized Stable Diffusion Models via Watermark LoRA
This repository hosts the official PyTorch implementation of the paper: "AquaLoRA: Toward White-box Protection for Customized Stable Diffusion Models via Watermark LoRA" (Accepted by ICML 2024).
🔍Method
We introduce AquaLoRA, the first implementation addressing white-box protection. Our approach integrates watermark information directly into the U-Net of SD models using a watermark LoRA module in a two-stage process. This module includes a scaling matrix, enabling flexible message updates without retraining. To maintain image fidelity, we design Prior Preserving Fine-Tuning to ensure watermark integration with small impact on the distribution. Our extensive experiments and ablation studies validate the effectiveness of our design.
🚀Getting Started
Prerequisites
Run the following command to install the environment:
git clone https://github.com/Georgefwt/AquaLoRA.git
cd AquaLoRA
conda create -n aqualora python=3.10
conda activate aqualora
pip install -r requirements.txt
Training
See train/README.md for details.
Evaluation
See evaluation/README.md for details.
Usage
We provide a script to merge LoRA into the Stable Diffusion model itself, enabling protection for customized models. You can run:
cd scripts
python diffusers_lora_to_webui.py --src_lora <your train folder>/<secret message>/pytorch_lora_weights.safetensors --tgt_lora watermark.safetensors
python merge_lora.py --sd_model <customized model safetensors> --models watermark.safetensors --save_to watermark_SDmodel.safetensors
Trained Models
| Stage | Link |
|---|---|
| Pretrained Latent Watermark | huggingface.co/georgefen/AquaLoRA-Models/tree/main/pretrained_latent_watermark |
| Prior-preserving Fine-tuned | huggingface.co/georgefen/AquaLoRA-Models/tree/main/ppft_trained |
| Robustness Enhanced (opt.) | huggingface.co/georgefen/AquaLoRA-Models/tree/main/rob_finetuned |
📝TODOs
- training code
- evaluation code
- release training dataset
- release trained checkpoints
🖥️Run Gradio Demo
Download AquaLoRA-Models from Hugging Face to the project's main directory, then run the Gradio demo by executing python run_gradio_demo.py.
🙏Acknowledgements
This code builds on the code from the diffusers library. In addition, we borrow code from the following repositories:
- sd-scripts for merge the LoRA with the model.
- dreamsim for calculate dreamsim metric.
- pytorch-fid for calculate FID metric.
📄Citation
If you use this code for your research, please cite the following work:
@article{feng2024aqualora,
title={AquaLoRA: Toward White-box Protection for Customized Stable Diffusion Models via Watermark LoRA},
author={Feng, Weitao and Zhou, Wenbo and He, Jiyan and Zhang, Jie and Wei, Tianyi and Li, Guanlin and Zhang, Tianwei and Zhang, Weiming and Yu, Nenghai},
journal={arXiv preprint arXiv:2405.11135},
year={2024}
}