算力平台:
将模型适应新任务
许多扩散系统共享相同的组件,允许你将为一项任务预训练的模型适应到完全不同的任务。
本指南将向你展示如何通过初始化和修改预训练的 [UNet2DConditionModel
] 的架构,将预训练的文本到图像模型适应到图像修复。
配置 UNet2DConditionModel 参数
[UNet2DConditionModel
] 默认情况下接受 输入样本 中的 4 个通道。例如,加载一个预训练的文本到图像模型,例如 stable-diffusion-v1-5/stable-diffusion-v1-5
,并查看 in_channels
的数量:
py
from diffusers import StableDiffusionPipeline
pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", use_safetensors=True)
pipeline.unet.config["in_channels"]
4
图像修复需要输入样本中的 9 个通道。你可以在一个预训练的图像修复模型中查看此值,例如 runwayml/stable-diffusion-inpainting
:
py
from diffusers import StableDiffusionPipeline
pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-inpainting", use_safetensors=True)
pipeline.unet.config["in_channels"]
9
为了将你的文本到图像模型用于图像修复,你需要将 in_channels
的数量从 4 更改为 9。
使用预训练的文本到图像模型权重初始化一个 [UNet2DConditionModel
],并将 in_channels
更改为 9。更改 in_channels
的数量意味着你需要设置 ignore_mismatched_sizes=True
和 low_cpu_mem_usage=False
以避免由于形状现在不同而导致的大小不匹配错误。
py
from diffusers import UNet2DConditionModel
model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
unet = UNet2DConditionModel.from_pretrained(
model_id,
subfolder="unet",
in_channels=9,
low_cpu_mem_usage=False,
ignore_mismatched_sizes=True,
use_safetensors=True,
)
文本到图像模型中其他组件的预训练权重从它们的检查点初始化,但 unet
的输入通道权重 (conv_in.weight
) 是随机初始化的。对模型进行微调以进行图像修复非常重要,否则模型会返回噪声。