モデルを新しいタスクに適応させる

多くの拡散モデルシステムは同じコンポーネントを共有しており、あるタスク用に事前学習されたモデルを、まったく異なるタスクに適応させることができます。

このガイドでは、事前学習済みのUNet2DConditionModelのアーキテクチャを初期化・修正することにより、事前学習済みのテキストから画像への変換モデルをインペインティングに適応させる方法を紹介します。

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チャネルが必要です。この値は、stable-diffusion-v1-5/stable-diffusion-inpaintingのような事前学習済みの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）はランダムに初期化されます。インペインティング用にモデルをファインチューニングすることが重要であり、そうしないとモデルはノイズを返します。