search

lthero-big / A-watermark-for-Diffusion-Models

This is an unofficial implementation of the Paper by Kejiang Chen et.al. on Gaussian Shading: Provable Performance-Lossless Image Watermarking for Diffusion Models
27 stars 2 forks source link

readme

A watermark for Diffusion Models

中文 | English

[!Note] This is an unofficial implementation of the Paper by Kejiang Chen et.al. on Gaussian Shading: Provable Performance-Lossless Image Watermarking for Diffusion Models

Features

Generated images

Watermarked image on the left | Image without watermark on the right

Images after distortions

Bit Accuracy Result

Tutorial for SD-CLI

Generating Watermarked Images

  1. Download and ensure the original Stable Diffusion project is able to generate images.
  2. Add the following codes into txt2img.py in the scripts folder of Stable Diffusion.

First, add the following codes at the end of the function parse_args()

parser.add_argument(
    "--message",
    type=str,
    default="",
    help="watermark message",
)
parser.add_argument(
    "--key_hex",
    type=str,
    default="5822ff9cce6772f714192f43863f6bad1bf54b78326973897e6b66c3186b77a7",
    help="key_hex",
)
parser.add_argument(
    "--nonce_hex",
    type=str,
    default="05072fd1c2265f6f2e2a4080a2bfbdd8",
    help="nonce_hex",
)

And you will get things like this

parser.add_argument(
    "--bf16",
    action='store_true',
    help="Use bfloat16",
)
# Here are the new codes
############################
parser.add_argument(
    "--message",
    type=str,
    default="",
    help="watermark message",
)
parser.add_argument(
    "--key_hex",
    type=str,
    default="5822ff9cce6772f714192f43863f6bad1bf54b78326973897e6b66c3186b77a7",
    help="key_hex",
)
parser.add_argument(
    "--nonce_hex",
    type=str,
    default="05072fd1c2265f6f2e2a4080a2bfbdd8",
    help="nonce_hex",
)
##############################
opt = parser.parse_args()
return opt

Second, add the following codes after the line for n in trange(opt.n_iter, *desc*="Sampling"):

from gs_insert import gs_watermark_init_noise
Z_s_T_arrays = [gs_watermark_init_noise(opt,opt.message) for _ in range(opt.n_samples)]
start_code = torch.stack([torch.tensor(Z_s_T_array).float() for Z_s_T_array in Z_s_T_arrays]).to(device)

It will be like

for n in trange(opt.n_iter, desc="Sampling"):
    # Here are the new codes
    ############################
    from gs_insert import gs_watermark_init_noise
    Z_s_T_arrays = [gs_watermark_init_noise(opt,opt.message) for _ in range(opt.n_samples)]
    start_code = torch.stack([torch.tensor(Z_s_T_array).float() for Z_s_T_array in Z_s_T_arrays]).to(device)
    ##############################
    for prompts in tqdm(data, desc="data"):
        uc = None
        if opt.scale != 1.0:
            uc = model.get_learned_conditioning(batch_size * [""])
  1. Execute the command below to generate watermarked images.
python scripts/txt2img.py --prompt "a professional photograph of an astronaut riding a horse" \
--ckpt ../ckpt/v2-1_512-ema-pruned.ckpt \
--config ./configs/stable-diffusion/v2-inference.yaml \
--H 512 --W 512  \
--device cuda \
--n_samples 2 \
--key_hex "5822ff9cce6772f714192f43863f6bad1bf54b78326973897e6b66c3186b77a7" \
--nonce_hex "" \
--message "lthero"

Parameter Explanation

[!important]

  • Both key_hex and nonce_hex can be left unentered, in which case a random 32-byte key_hex and a random 16-byte nonce_hex will be generated.
  • The message can also be left blank, in which case a random 256-bit (32-byte) content will be generated.
  • All the above parameters will be saved in info_data.txt (located in the root directory of the Stable Diffusion project).
  • If it is the first run, just keep key_hex and nonce_hex be unentered, so that the code is automatically generated; Or use the following code to generate key_hex and nonce_hex
import os
key = os.urandom(32)
nonce = os.urandom(16)
print(key.hex())
print(nonce.hex())

[Recommended]Tutorial for SD-WebUI

This work implement watermark embedding functionality in the form of a script based on the Stable Diffusion-WebUI project

Script Installation

  1. Place the GS_watermark_insert.py file from this project's scripts directory into the scripts directory of Stable Diffusion-WebUI.
  2. After you restart webui, you will find the script options at the bottom of the txt2img and img2img sections
  3. Click the script options then you can use "GS_watermark_insert".

image-script in webui

Three Parameters Provided by the Script

[!note]

  • You can only fill in the Key and leave the Nonce empty; it will automatically select a Nonce.
  • Both Key and Nonce can be left blank, in this case, Key and Nonce will be generated automatically.

[!important]

In the root directory of Stable Diffusion-WebUI, you can find info_data.txt, which records the Key, Nonce, and Message.

Image Generation

Extracting Watermark Messages

Method 1

  1. Modify the parameters inside extricate.py.
  2. Run extricate.py with the command python extricate.py.

Method 2

Pass parameters through the command line:

python extricate.py \
--single_image_path "path to image" \
--image_directory_path "directory to image" \
--key_hex "xxxxxxxxxx" \
--original_message_hex "xxxxxxxxxxxxx" \
--num_inference_steps 50 \
--scheduler "DDIM" \
--is_traverse_subdirectories 0

Parameter Explanation

[!caution]

The original_message_hex must be input in hexadecimal format, strictly following what is entered in info_data.txt.

After running extricate.py, it will output the image name and Bit accuracy:

v2-1_512_00098-3367722000JPEG_QF_95.jpg
Bit accuracy:  1.0

[!note]

If batch processing is used, a result.txt file will be generated in the input directory, recording the results for each image.

If recursive processing is used, each subdirectory under image_directory_path will have a result.txt file, and image_directory_path will have a result.txt recording the average Bit accuracy in each subdirectory.