Diffusers documentation

将模型适配到新任务

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

将模型适配到新任务

许多扩散系统共享相同的组件,因此你可以将一个针对某项任务预训练的模型适配到完全不同的任务上。

本指南将展示如何通过初始化并修改预训练 UNet2DConditionModel 的架构,将预训练的文生图模型适配用于图像修复(inpainting)。

配置 UNet2DConditionModel 参数

默认情况下,UNet2DConditionModel 接受 输入样本 中的 4 个通道。例如,加载一个预训练的文生图模型,如 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

Inpainting 要求输入样本具有 9 个通道。您可以在预训练的 inpainting 模型(例如 stable-diffusion-v1-5/stable-diffusion-inpainting)中检查此值:

from diffusers import StableDiffusionPipeline

pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/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)是随机初始化的。对模型进行修复任务的微调非常重要,否则模型会返回噪声。

Update on GitHub