我们主要想将该项目用在stable diffusion上,所以要下下载好stable diffusion,该项目的复现教程较多,此处不过多赘述.
1.克隆Stable Diffusion项目
git clone https://github.com/CompVis/stable-diffusion-public cd stable-diffusion-public
下载模型文件并放在models中
2.在本地克隆stable-signature
git clone https://github.com/facebookresearch/stable_signature cd stable_signature
安装相关依赖
conda install -c pytorch torchvision pytorch==1.12.0 cudatoolkit==11.3 pip install -r requirements.txt
3.下载水印提取器
mkdir models wget https://dl.fbaipublicfiles.com/ssl_watermarking/dec_48b_whit.torchscript.pt -P models/ wget https://dl.fbaipublicfiles.com/ssl_watermarking/other_dec_48b_whit.torchscript.pt -P models/
由于我们使用的是stable diffusion v1,所以不需要修改检查点.
4.权重下载
git clone https://github.com/SteffenCzolbe/PerceptualSimilarity.git cp -r PerceptualSimilarity/src/loss/weights src/loss/losses/ rm -r PerceptualSimilarity
5.修改源代码中text2img文件
找到加载模型的代码块,在加载模型的代码后面插入以下代码,以加载Stable Signature的水印解码器权重:
点击查看代码
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")