WatermarkAttacker

[Update] The paper is accepted by NeurIPS 2024!

Welcome to contribute to this project!

This repository contains the code for the paper Invisible Image Watermarks Are Provably Removable Using Generative AI.

If you find this repository useful, please cite our paper:

@article{zhao2023invisible,
  title={Invisible Image Watermarks Are Provably Removable Using Generative AI},
  author={Zhao, Xuandong and Zhang, Kexun and Su, Zihao and Vasan, Saastha and Grishchenko, Ilya and Kruegel, Christopher and Vigna, Giovanni and Wang, Yu-Xiang and Li, Lei},
  journal={arXiv preprint arXiv:2306.01953},
  year={2023}
}

We propose a family of regeneration attacks to remove invisible image watermarks. The attack method effectively removes invisible watermarks.

Our attack first maps the watermarked image to its embedding, which is another representation of the image. Then the embedding is noised to destruct the watermark. After that, a regeneration algorithm reconstructs the image from the noisy embedding. As shown in the figure below:

Example

First, install the dependencies:

pip install -r requirements.txt

Then you can try the demo in demo.ipynb.

For Regen-Diffusion, you can set the DiffWMAttacker noise_step=30/60/100... to control the noise level. The larger the noise_step, the more noise is added to the image, and the watermark is more likely to be removed.

For Regen-VAE, you can set the VAEWMAttacker quality=6/5/4/3/2/1 to control the noise level. The smaller the quality, the more noise is added to the image, and the watermark is more likely to be removed.

This demo attacks invisible-watermark which is used in stable diffusion.

How it works?

In short, our method can be described using the following equation (1)

$$ \begin{equation} \hat{x} = \underbrace{\mathcal{A}\Big( \overbrace{\phi(x_w) + \mathcal{N}(0,\sigma^2 I_d)}^{\text{destructive}}\Big)}_{\text{constructive}}, \end{equation} $$

where $x_w$ is the watermarked image, $\phi$ is an embedding function, $\mathcal{N}$ is a Gaussian distribution that adds noise to the embedding, and $\mathcal{A}$ is a reconstruction algorithm that takes the noisy embedding of the image and tries to construct the original image.

The embedding function $\phi$ and the reconstruction algorithm $\mathcal{A}$ can be instantiated with many algorithms.

Here let's look at how Stable Diffusion can instantiate our algorithm (see this blog for more information on Stable Diffusion).

The above figure shows the architecture of Stable Diffusion.

The destructive process in our algorithm corresponds to adding noise to the image until a certain time step $t$. To do this, Stable Diffusion first maps the image from the pixel space to a latent space using its VAE encoder $\mathcal{E}$, and then adds noise to the latent space using a diffusion process, which is defined in the following equation (2)

$$ \begin{equation} z_t=\overbrace{\sqrt{\alpha(t)}\mathcal{E}(x_w)}^{\phi(x_w)}+\mathcal{N}(0,\overbrace{(1-\alpha(t))}^{\sigma^2}I_d), \end{equation} $$

where $\alpha(t)$ is a function that controls the amount of noise added at time step $t$. By comparing equations (1) and (2), we can see that the embedding function $\phi$ is the scaled VAE encoder $\sqrt{\alpha(t)}\mathcal{E}$.

After obtaining the noised representation $z_t$, we denoise it using Stable Diffusion to reconstruct the image, which corresponds to $\mathcal{A}$ in equation (1).