树轮水印：扩散图像的隐形且稳健的指纹

<img src=scripts/teaser.png width="80%" height="60%">

这是树轮水印的官方实现代码。

如有任何问题，请随时发邮件给Yuxin（ywen@umd.edu）。

关于

我们提出了树轮水印技术来为扩散模型输出添加水印。树轮水印通过选择初始噪声数组，使其傅里叶变换在中心附近包含一个精心构造的模式。这个模式被称为密钥。然后使用标准扩散流程将这个初始噪声向量转换为图像，无需任何修改。要检测图像中的水印，需要反转扩散模型以检索用于生成的原始噪声数组。然后检查该数组是否包含密钥。

依赖

• PyTorch == 1.13.0
• transformers == 4.23.1
• diffusers == 0.11.1
• datasets

注意：更高版本的diffusers可能与DDIM反转代码不兼容。

使用方法

执行主要实验并计算CLIP评分

对于非对抗性情况，你可以直接运行：

python run_tree_ring_watermark.py --run_name no_attack --w_channel 3 --w_pattern ring --start 0 --end 1000 --with_tracking --reference_model ViT-g-14 --reference_model_pretrain laion2b_s12b_b42k

你可以修改参数来执行攻击。例如，对于75度旋转：

python run_tree_ring_watermark.py --run_name rotation --w_channel 3 --w_pattern ring --r_degree 75 --start 0 --end 1000 --with_tracking

对于更多对抗性情况，请参见这里。

关于论文中提到的其他水印类型，你可以查看scripts/。

计算FID

你可以在这里下载论文中使用的5000个COCO示例。你可以根据fid_outputs/coco/meta_data.json的格式自由添加更多数据或其他数据集。

然后，要计算FID，你可以运行：

python run_tree_ring_watermark_fid.py --run_name fid_run --w_channel 3 --w_pattern ring --start 0 --end 5000 --with_tracking --run_no_w

对ImageNet模型执行主要实验

你可以在这里获取预训练模型。例如，论文使用的模型链接是256x256扩散。然后，你可以按照脚本运行实验。

参数

树轮水印的关键超参数：

• w_channel：水印通道的索引。如果设为-1，则对所有通道添加水印。
• w_pattern：水印类型：zeros、rand、ring。
• w_radius：水印半径。

欢迎提出建议和拉取请求！