使模型适应新任务

许多扩散系统共享相同的组件，允许您将一个预训练模型适应于完全不同的任务。

本指南将向您展示如何通过初始化和修改预训练的 UNet2DConditionModel 架构，将预训练的文本到图像模型适应于图像修复。

配置 UNet2DConditionModel 参数

UNet2DConditionModel 默认接受输入样本中的 4 个通道（input sample）。例如，加载一个预训练的文本到图像模型，如 stable-diffusion-v1-5/stable-diffusion-v1-5，并查看 in_channels 的数量。

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。

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 以避免由于形状不同而导致的尺寸不匹配错误。

from diffusers import AutoModel

model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
unet = AutoModel.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）是随机初始化的。对图像修复模型进行微调很重要，否则模型会返回噪声。

< > 在 GitHub 上更新

Diffusers

使模型适应新任务

配置 UNet2DConditionModel 参数