Diffusers 文档

将模型适配到新任务

Hugging Face's logo
加入 Hugging Face 社区

并获得增强文档体验

开始使用

将模型适配到新任务

许多扩散系统共享相同的组件,使您可以将为一项任务预训练的模型适配到完全不同的任务。

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

配置 UNet2DConditionModel 参数

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

from diffusers import StableDiffusionPipeline

pipeline = StableDiffusionPipeline.from_pretrained("runwayml/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=Truelow_cpu_mem_usage=False 以避免大小不匹配错误,因为形状现在不同了。

from diffusers import UNet2DConditionModel

model_id = "runwayml/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) 是随机初始化的。对模型进行修复任务的微调非常重要,否则模型会返回噪声。

< > 在 GitHub 上更新