Diffusers代码学习-ControlNet(图生图）

之前写过通过姿态图片文生图（见这里Diffusers代码学习-ControlNet）

对于图生图，通常会将初始图像和提示词传递给管道以生成新图像。使用ControlNet，您可以传递额外的条件输入来指导模型。让我们用深度图来调节模型，深度图是一个包含空间信息的图像。这样，ControlNet可以使用深度图作为控件来引导模型生成保留空间信息的图像。

以下将使用[StableDiffusionControlNetImg2ImgPipeline]执行此任务，这与[StableDiffusionControlNetPipeline]不同，因为它允许传递初始图像作为图像生成过程的起点。同时加载图像并使用的深度估计[~transformer.Pipeline] 提取图像深度图的转换器。

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

import torch
import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
import numpy as np
from transformers import pipeline
from diffusers.utils import load_image, make_image_grid

# 以下代码引入图像处理环节，并提取图像深度

image = load_image(
"https://hf-mirror.com/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet-img2img.jpg"
)

def get_depth_map(image, depth_estimator):
image = depth_estimator(image)["depth"]
image = np.array(image)
image = image[:, :, None]
image = np.concatenate([image, image, image], axis=2)
detected_map = torch.from_numpy(image).float() / 255.0
depth_map = detected_map.permute(2, 0, 1)
return depth_map

depth_estimator = pipeline("depth-estimation")
depth_map = get_depth_map(image, depth_estimator).unsqueeze(0).half().to("cuda")

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

from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, UniPCMultistepScheduler

controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11f1p_sd15_depth", torch_dtype=torch.float16, use_safetensors=True)

pipe = StableDiffusionControlNetImg2ImgPipeline.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(
"lego batman and robin", image=image, control_image=depth_map,
).images[0]

make_image_grid([image, output], rows=1, cols=2)

以下为提供深度信息的参考图片

以下为输出结果

AIGC应用37