Diffusers代码学习-ControlNet(Inpaint）

对于Inpaint，需要一个初始图像、一个蒙版图像和一个描述用什么替换蒙版的提示词。ControlNet模型允许添加另一个控制图片来调节模型。让我们用Inpaint蒙版来调整模型。这样，ControlNet可以使用修复掩模作为控件来引导模型在蒙版区域内生成图像。

# 以下代码为程序运行进行设置

import os

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
from diffusers.utils import load_image, make_image_grid

# 以下代码加载初始图像，并调整图像大小

init_image = load_image(
    "https://hf-mirror.com/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet-inpaint.jpg"
)
init_image = init_image.resize((512, 512))

# 以下代码加载蒙版图像，并调整图像大小

mask_image = load_image(
    "https://hf-mirror.com/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet-inpaint-mask.jpg"
)
mask_image.show()
mask_image = mask_image.resize((512, 512))


import numpy as np
import torch

# 创建一个函数，从初始图像和蒙版图像中准备控制图像。

# 这将创建一个张量，如果蒙版图像中的相应像素超过某个阈值，则将原始图像中的像素标记为掩码。

def make_inpaint_condition(image, image_mask):
image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0

assert image.shape[0:1] == image_mask.shape[0:1]
image[image_mask > 0.5] = -1.0 # set as masked pixel
image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)
image = torch.from_numpy(image)
return image

control_image = make_inpaint_condition(init_image, mask_image)

# 加载以Inpaint为条件的ControlNet模型，并将其传递给[SableDiffusionControlNetInpaintPipeline]。

# 使用更快的[UniPCMultipstepScheduler]并启用模型卸载机制，以加快推理并减少内存使用。

from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, UniPCMultistepScheduler

# 以下代码会引入与ControlNet相关的模型，及Scheduler

controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16,
                                             use_safetensors=True)
pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True
)

pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)

pipe.enable_model_cpu_offload()

# 以下代码由提示词及原始图片，蒙版图片形成控制后生成图片

output = pipe(
    "corgi face with large ears, detailed, pixar, animated, disney",
    num_inference_steps=20,
    eta=1.0,
    image=init_image,
    mask_image=mask_image,
    control_image=control_image,
).images[0]

以下为原始图像