Stable Signature运行流程

先安装原始的stable diffusion，目前stable diffusion有两个版本(v1,v2)，建议用v2（v1实际跑越来的问题太多了）；

除了原始的sd，目前还有stable diffusion webui版本，这个版本拥有了UI，而且一键安装不需要自己配置环境，上手即用，但缺点是，目前无法对webui中的txt2img.py文件进行修改，即：无法将stable signature嵌入进去。

stable signature是基于stable diffusion基础上的，所以，要先运行好stable diffusion

下面是stable diffusion【v2】安装流程

stable diffusion【v2】项目：https://github.com/Stability-AI/stablediffusion

注意：作者是Stability-AI

下载stable diffusion项目

打开项目的目录，里面有environment.yaml文件，可以打开environment.yaml文件，对name进行修改（默认为ldm）

安装初始环境

随后，使用anaconda的shell，进入到项目的目录，并执行下面命令

1
2
conda env create -f environment.yaml
conda activate ldm

安装必要的库

随后，再执行下面的命令

1
2
3
conda install pytorch==1.12.1 torchvision==0.13.1 -c pytorch
pip install transformers==4.19.2 diffusers invisible-watermark
pip install -e .

至此，就可以运行stable diffusion了，但生成速度很慢，需要再安装xformers（交叉注意力相关），项目：https://github.com/facebookresearch/xformers

安装xformers的命令

1
conda install xformers -c xformers
  • (RECOMMENDED, linux & win) Install latest stable with pip: Requires PyTorch 2.1.0
1
2
3
4
# cuda 11.8 version
pip3 install -U xformers --index-url https://download.pytorch.org/whl/cu118
# cuda 12.1 version
pip3 install -U xformers --index-url https://download.pytorch.org/whl/cu121

运行stable diffusion项目

在项目的目录中，执行下面命令

1
2
3
4
5
# 成功出图
python scripts/txt2img.py --prompt "a professional photograph of an astronaut riding a horse" --ckpt ..\ckpt\512-ema-pruned.ckpt --config .\configs\stable-diffusion\v2-inference.yaml --H 512 --W 512  --device cuda --n_samples 1
---------------
# 或者
python scripts/txt2img.py --prompt "a professional photograph of an astronaut riding a horse" --ckpt 768-ema-pruned.ckpt --config configs/stable-diffusion/v2-inference-v.yaml --H 768 --W 768 --device cuda

参数解读

–prompt

提示词

–device cuda

强制使用cude，一定要添加，不然报错

–ckpt

是图片模型文件（weight）

ckpt文件的下载链接：https://huggingface.co/stabilityai/stable-diffusion-2-1/tree/main

官网提供了ckpt和safetensor两种版本，目前两个版本可以相互转换

相互转换的工具链接：https://github.com/diStyApps/Safe-and-Stable-Ckpt2Safetensors-Conversion-Tool-GUI

–config

配置文件

https://github.com/Stability-AI/stablediffusion/tree/main/configs/stable-diffusion

ckpt和config文件要相互对应，否则画出来的东西很抽象

如

模型：768-ema-pruned.ckpt要对应配置文件v2-1_768-ema-pruned.yaml
模型：512-ema-pruned.ckpt要对应配置文件v2-inference.yaml

在本命令中，使用的是下面两种相对应的，这里的v2-inference-v.yaml就是v2-1_768-ema-pruned.yaml，两个内容是一样的，所以还是按照模型和配置文件对应的定理
768model.ckpt --config configs/stable-diffusion/v2-inference-v.yaml

–H

高度，同理–W是宽度

–n_samples

批次数，默认设置3，会生成9张图，若设置5则生成15张图。

–seed

种子，默认是42，如果执行两次命令，使用相同的seed，生成出来的图片是相同的

在outputs\txt2img-samples 可以看到生成的图片，下面将使用stable signature

stable signature运行流程

确保stable diffusion能正常运行后

stable signature项目地址：https://github.com/facebookresearch/stable_signature

项目中提供了一个已经训练好的decoder，下载这个decoder即可使用—> 下载链接

使用decoder

需要将LDM自带的decoder切换到“下载的decoder”，需要做的

  1. 打开stable diffusion的txt2img.py文件（它的路径stablediffusion\scripts\txt2img.py）
  2. 在220行左右添加下面的代码
  3. 还要注释掉 img = put_watermark(img, wm_encoder) 这行代码，不让stable diffusion自带的水印函数生效

Reload weights of the LDM decoder in the Stable Diffusion scripts by appending the following lines after loading the checkpoint (for instance, L220 in the SD repo)

1
2
3
4
state_dict = torch.load(path/to/ldm/checkpoint_000.pth)['ldm_decoder']
msg = model.first_stage_model.load_state_dict(state_dict, strict=False)
print(f"loaded LDM decoder state_dict with message\n{msg}")
print("you should check that the decoder keys are correctly matched")

大致添加的位置如下

1
2
3
4
5
6
7
8
9
10
11
12
def main(opt):
    seed_everything(opt.seed)

    config = OmegaConf.load(f"{opt.config}")
    device = torch.device("cuda") if opt.device == "cuda" else torch.device("cpu")
    model = load_model_from_config(config, f"{opt.ckpt}", device)
    
    # 下面是添加的代码
    state_dict = torch.load(path/to/ldm/checkpoint_000.pth)['ldm_decoder']
    msg = model.first_stage_model.load_state_dict(state_dict, strict=False)
    print(f"loaded LDM decoder state_dict with message\n{msg}")
    print("you should check that the decoder keys are correctly matched")

如果使用**“下载的decoder”（生成此decoder的命令），而不是使用自己训练的decoder，那么就不需要选择ldm_decoder**这个键，即state_dict = torch.load(path/to/ckpt.pth)，在220行左右添加下面的代码，并且仍需要注释掉 img = put_watermark(img, wm_encoder) 这行代码

1
2
3
4
state_dict = torch.load(path/to/ldm/checkpoint_000.pth)
msg = model.first_stage_model.load_state_dict(state_dict, strict=False)
print(f"loaded LDM decoder state_dict with message\n{msg}")
print("you should check that the decoder keys are correctly matched")

随后，按上面运行stable diffusion的命令执行即可生成包含水印的图片

自己训练decoder

下载Data集

The paper uses the COCO dataset to fine-tune the LDM decoder (we filtered images containing people). All you need is around 500 images for training (preferably over 256x256).

下载LDM configs and checkpoints

Create LDM configs and checkpoints from the Hugging Face and Stable Diffusion repositories. The code should also work for Stable Diffusion v1 without any change. For other models (like old LDMs or VQGANs), you may need to adapt the code to load the checkpoints.

这里的configs和ckeckpoints要对应好

下载提取器文件

想提取水印，需要有个提取器文件，即dec_48b_whit.torchscript.pt，下载方式如下

pth文件没有经过白化处理，pt文件经过了处理可以直接用

dec_48b_whit.torchscript.pt 是可以直接用的

other_dec_48b_whit.torchscript.pt 鲁棒性可能更好，但会稍微降低图片质量

Model Checkpoint Torch-Script
Extractor dec_48b.pth dec_48b_whit.torchscript.pt
Other other_dec_48b_whit.pth other_dec_48b_whit.torchscript.pt

下载Perceptual Losses

The perceptual losses are based on this repo. You should download the weights here: https://github.com/SteffenCzolbe/PerceptualSimilarity/tree/master/src/loss/weights, and put them in a folder called losses (this is used in src/loss/loss_provider.py#L22). To do so you can run

1
2
3
git clone https://github.com/SteffenCzolbe/PerceptualSimilarity.git
cp -r PerceptualSimilarity/src/loss/weights src/loss/losses/
rm -r PerceptualSimilarity

生成微调解码器decoder

下载完上面的东西，现在可以使用以下命令来生成微调的decoder

1
2
3
4
5
6
python finetune_ldm_decoder.py --num_keys 1 \
    --ldm_config path/to/ldm/config.yaml \
    --ldm_ckpt path/to/ldm/ckpt.pth \
    --msg_decoder_path path/to/msg/decoder/ckpt.torchscript.pt \
    --train_dir path/to/train/dir \
    --val_dir path/to/val/dir

参数解读

  1. ldm_config 是ldm配置文件，一般叫"v2-inference.ymal"
  2. ldm_ckpt 是ldm的模型权重，一般为ckpt格式
  3. msg_decoder_path 是HiDDen的提取器，即dec_48b_whit.torchscript.pt
  4. train_dir和val_dir是训练集和验证集目录

执行后会生成下面三种文件

  • num_keys checkpoints of the LDM decoder with watermark fine-tuning (checkpoint_000.pth, etc.),
  • keys.txt: text file containing the keys used for fine-tuning (one key per line),
  • imgs: folder containing examples of auto-encoded images.

Params of LDM fine-tuning used in the paper
Logs during LDM fine-tuning

提取水印

打开decoding.ipynb文件，把dec_48b_whit.torchscript.pt路径填写进去

1
2
msg_extractor = torch.jit.load("E:/C_code/Python/AIGC/stable_signature/hidden/models/dec_48b_whit.torchscript.pt").to("cuda")

顺便把生成好的包含水印图片的路径写进去

1
img = Image.open("E:/C_code/Python/AIGC/stablediffusion_v2/outputs/txt2img-samples/samples/batch_02_withmark/qq11.png")

随后按顺便运行decoding.ipynb文件即可提取出key，再与decoding.ipynb文件中的key对比，如果bit accuracy比较高（大于0.8，基本则说明包含水印），对于一个不包含水印的图片，bit accuracy一般在0.5浮动

鲁棒性测试

The run_eval.py script can be used to get the robustness and quality metrics on a folder of images. For instance:

1
2
3
python run_eval.py --eval_imgs False --eval_bits True \
    --img_dir path/to/imgs_w \
    --key_str '111010110101000001010111010011010100010000100111'

will return a csv file containing bit accuracy for different attacks applied before decoding.

1
2
python run_eval.py --eval_imgs True --eval_bits False \
    --img_dir path/to/imgs_w --img_dir_nw path/to/imgs_nw

will return a csv file containing image metrics (PSNR, SSIM, LPIPS) between watermarked (_w) and non-watermarked (_nw) images.

运行时遇到的一些问题

问题一：画出的东西很奇怪，很抽像

https://www.reddit.com/r/StableDiffusion/comments/zlzn70/i_switched_to_v21_768nonemapruned_and_now/

来自reddit的回复

You can download YAML files that work with v2-1 checkpoints from this webpage:

https://github.com/Stability-AI/stablediffusion/tree/main/configs/stable-diffusion

Use the v2-inference.yaml with the v2-1_512-ema-pruned.ckpt and v2-1_512-nonema-pruned.ckpt files.

Use the v2-inference-v.yaml with the v2-1_768-ema-pruned.ckpt and v2-1_768-nonema-pruned.ckpt files.

In order to properly use a YAML file, copy the appropriate YAML file to the same directory as the appropriate CKPT file. Then you rename the YAML so that it has the exact same name as corresponding CKPT file. Do not change the extensions of the file names.

问题二：无法下载checkpoint_liberty_with_aug.pth

把这个项目下载下来：https://github.com/DagnyT/hardnet/tree/master，在这个项目的hardnet/pretrained/train_liberty_with_aug/路径有个checkpoint_liberty_with_aug.pth文件

把checkpoint_liberty_with_aug.pth粘贴到下面的目录即可
C:\Users\12904.cache\torch\hub\checkpoints\checkpoint_liberty_with_aug.pth

原因：SD项目里面的下载链接无效了，导致这个文件无法被下载

问题三：不使用–device cuda报错

具体报错不记得了

将safetenors转成ckpt

项目：https://github.com/huggingface/diffusers/tree/main/scripts
工具：https://github.com/diStyApps/Safe-and-Stable-Ckpt2Safetensors-Conversion-Tool-GUI