基于 huggingface diffuser 库云服务器实现 stable diffusion inpaint样例代码

---

一、stable diffusion inpant样例

（一）理论基础

Stable Diffusion是计算机视觉领域的一个生成式大模型，可以用于文生图，图生图，图像inpainting，ControlNet控制生成，图像超分等丰富的任务。inpaint是Stable Diffusion仅重绘图像部分的技术，将画面中被手工遮罩的部分重新绘制，使用stable diffusion实现inpaint的理论可以参考我之前的文章：Stable Diffusion原理解析-inpaint修复图片。

（二）样例背景

huggingface收纳了许多最前沿的模型和数据集等有趣的工作，其中包括了diffusers库，基于 huggingface diffuser 库可以自行在云服务器实现一系列有关操作。对于stable diffusion inpaint，huggingface提供了一个运行样例，大致内容如下：

二、本地部署

接下来实战一下本地部署。

（一）环境配置

1. 创建一个diffenv环境用于实现上述内容

conda create -n diffenv python=3.8
conda activate diffenv
1
2

2. 下载相关库

pip install diffusers==0.4.0
pip install transformers scipy ftfy
pip instal1 torch torchvision torchaudio
pip install Pillow
pip install requests
pip install --upgrade diffusers[torch]
1
2
3
4
5
6

（二）模型下载

有两种方式将huggingface中的stable-diffusion-inpainting模型下载到本地：

• 方式一：采用git远程clone下来，具体方式可以参考之前的内容：huggingface学习 | 云服务器使用git-lfs下载huggingface上的模型文件；
• 方式二：在官网上直接下载，注意不要有遗漏，下载后的文件储存地址要与官网上的地址类似。
  

（三）参照样例编写代码

1. 导入相关库

import PIL
import requests
import torch
from io import BytesIO

from diffusers import StableDiffusionInpaintPipeline
1
2
3
4
5
6

2. 加载模型

#加载Stable Diffusion Inpainting模型并创建一个可以用于图像修复的pipeline
pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "./stable-diffusion-inpainting",#下载的模型所在的本地地址
    revision="fp16",#使用FP16（半精度浮点数）进行训练
    torch_dtype=torch.float16)


pipe = pipe.to("cuda")#在GPU上进行模型推断
1
2
3
4
5
6
7
8

3. 读取本地图像内容（模型要求图片格式为PIL）

img_path="./example.png"
mask_path="./example-mask.png"

#按照路径打开文件并读取字节数据，然后将字节数据传递给BytesIO，最后用PIL库打开并处理图像
def download_image(path):
    with open(path,'rb') as file:
        image_data=file.read()
    return PIL.Image.open(BytesIO(image_data)).convert("RGB")

init_image = download_image(img_path).resize((512, 512))
mask_image = download_image(mask_path).resize((512, 512))

1
2
3
4
5
6
7
8
9
10
11
12

4. 模型运行

prompt = "Face of a yellow cat, high resolution, sitting on a park bench"
image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]
image.save("./yellow_cat_on_park_bench.png")
1
2
3

5. 实现效果
   输入"Face of a yellow cat, high resolution, sitting on a park bench"运行之后为：
   

三、全部代码

import PIL
import requests
import torch
from io import BytesIO

from diffusers import StableDiffusionInpaintPipeline

'''
def download_image(url):
    response = requests.get(url)
    return PIL.Image.open(BytesIO(response.content)).convert("RGB")


img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"

init_image = download_image(img_url).resize((512, 512))
mask_image = download_image(mask_url).resize((512, 512))
'''

img_path="./example.png"
mask_path="./example-mask.png"
def download_image(path):
    with open(path,'rb') as file:
        image_data=file.read()
    return PIL.Image.open(BytesIO(image_data)).convert("RGB")

init_image = download_image(img_path).resize((512, 512))
mask_image = download_image(mask_path).resize((512, 512))

pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "./stable-diffusion-inpainting",revision="fp16",torch_dtype=torch.float16)


pipe = pipe.to("cuda")

prompt = "Face of a yellow cat, high resolution, sitting on a park bench"
image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]
image.save("./yellow_cat_on_park_bench.png")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39