Custom Diffusion

Custom Diffusionは、画像生成モデルをパーソナライズするための学習手法です。Textual Inversion、DreamBooth、LoRAと同様に、Custom Diffusionにはわずか数枚（約4〜5枚）のサンプル画像のみが必要です。この手法は、クロスアテンション層の重みのみを学習することで機能し、新たに学習した概念を表すために特別な単語を使用します。Custom Diffusionは、同時に複数の概念を学習できる点で独自性を持っています。

vRAM が限られた GPU で学習を行う場合は、--enable_xformers_memory_efficient_attention を有効にして xFormers を使用することをお勧めします。これにより、vRAM 要件（16GB）を下げながら高速な学習が可能になります。さらにメモリを節約するには、学習引数に --set_grads_to_none を追加し、勾配をゼロではなく None に設定します（このオプションは問題を引き起こす可能性があるため、問題が発生した場合はこのパラメータを削除してみてください）。

このガイドでは、train_custom_diffusion.py スクリプトを詳しく見ていき、その仕組みをより深く理解したうえで、ご自身のユースケースに合わせてどのようにカスタマイズできるかを解説します。

スクリプトを実行する前に、ソースからライブラリをインストールしていることを確認してください。

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

学習スクリプトを含むサンプルフォルダに移動し、必要な依存関係をインストールします:

cd examples/custom_diffusion
pip install -r requirements.txt
pip install clip-retrieval

🤗 Accelerate は、複数 GPU/TPU での学習や混合精度学習を支援するためのライブラリです。ハードウェアや環境に基づいて、学習環境を自動的に設定します。詳細については、🤗 Accelerate のクイックツアーをご覧ください。

🤗 Accelerate の環境を初期化します:

accelerate config

設定を選択せずにデフォルトの🤗 Accelerate環境を設定するには:

accelerate config default

または、ノートブックのようにインタラクティブシェルをサポートしていない環境では、以下を使用できます。

from accelerate.utils import write_basic_config

write_basic_config()

最後に、独自のデータセットでモデルをトレーニングしたい場合は、トレーニング用のデータセットを作成するガイドを参照して、トレーニングスクリプトと互換性のあるデータセットの作成方法を確認してください。

以下のセクションでは、スクリプトの修正方法を理解するために重要な部分に焦点を当てていますが、スクリプトのすべての側面を詳細に説明するものではありません。さらに詳しく知りたい場合は、スクリプトをぜひお読みいただき、ご質問やご懸念事項があればお知らせください。

スクリプトのパラメータ

トレーニングスクリプトには、トレーニング実行をカスタマイズするためのすべてのパラメータが含まれています。これらは parse_args() 関数内で定義されています。この関数にはデフォルト値が設定されていますが、必要に応じてトレーニングコマンド内で独自の値を設定することもできます。

例えば、入力画像の解像度を変更するには:

accelerate launch train_custom_diffusion.py \
  --resolution=256

多くの基本的なパラメータはDreamBoothのトレーニングガイドで説明されているため、このガイドではCustom Diffusion固有のパラメータに焦点を当てます:

--freeze_model: クロスアテンション層のキーとバリューのパラメータを凍結します。デフォルトはcrossattn_kvですが、crossattnに設定すると、クロスアテンション層のパラメータをすべて学習できます。
--concepts_list: 複数の概念を学習するには、概念を含む JSON ファイルへのパスを指定します
--modifier_token: 学習された概念を表すために使用される特別な単語
--initializer_token: modifier_tokenの埋め込みを初期化するために使用される特別な単語

事前分布の保存損失

Prior preservation loss は、モデル自身が生成したサンプルを用いて、より多様な画像の生成方法を学習させる手法です。これらの生成されたサンプル画像は、ユーザーが提供した画像と同じクラスに属するため、モデルがそのクラスについて学習した内容を保持し、クラスに関する既存の知識を活用して新しい構成を作成するのを助けます。

事前保持損失のパラメータの多くは、DreamBoothのトレーニングガイドで説明されています。

正則化

Custom Diffusionでは、過学習を防ぐために、実写の小さな画像セットでターゲット画像を学習します。わずか数枚の画像でのみ学習を行う場合、これは容易に実行できることが想像できるでしょう。clip_retrievalを使用して200枚の実写画像をダウンロードします。class_promptは、ターゲット画像と同じカテゴリである必要があります。これらの画像はclass_data_dirに保存されます。

python retrieve.py --class_prompt cat --class_data_dir real_reg/samples_cat --num_class_images 200

正則化を有効にするには、次のパラメータを追加します。

--with_prior_preservation: 事前分布保持損失を使用するかどうか
--prior_loss_weight: モデルに対する事前分布保持損失の影響度を制御します
--real_prior: 過学習を防ぐために、実際の画像の小さなセットを使用するかどうか

accelerate launch train_custom_diffusion.py \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --class_data_dir="./real_reg/samples_cat" \
  --class_prompt="cat" \
  --real_prior=True \

学習スクリプト

Custom Diffusionの学習スクリプトには、DreamBoothスクリプトと類似したコードが多く含まれています。本ガイドでは、Custom Diffusionに関連するコードに焦点を当てて説明します。

Custom Diffusionの学習スクリプトには、2つのデータセットクラスがあります。

CustomDiffusionDataset: 学習用に画像、クラス画像、プロンプトを前処理します
PromptDataset: クラス画像の生成用にプロンプトを準備します

次に、modifier_token がトークナイザーに追加され、トークン ID に変換されます。その後、新しい modifier_token を考慮して、トークン埋め込みのサイズが調整されます。続いて、modifier_token の埋め込みは、initializer_token の埋め込みで初期化されます。テキストエンコーダーのパラメータはすべて凍結されますが、モデルが概念に関連付けようとしているのがトークン埋め込みであるため、トークン埋め込みのみは凍結されません。

params_to_freeze = itertools.chain(
    text_encoder.text_model.encoder.parameters(),
    text_encoder.text_model.final_layer_norm.parameters(),
    text_encoder.text_model.embeddings.position_embedding.parameters(),
)
freeze_params(params_to_freeze)

次に、Custom Diffusion weightsをattentionレイヤーに追加する必要があります。これは、attention weightsのshapeとsizeを正しく設定し、各UNetブロックに適切な数のattention processorsを設定するために非常に重要なステップです。

st = unet.state_dict()
for name, _ in unet.attn_processors.items():
    cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
    if name.startswith("mid_block"):
        hidden_size = unet.config.block_out_channels[-1]
    elif name.startswith("up_blocks"):
        block_id = int(name[len("up_blocks.")])
        hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
    elif name.startswith("down_blocks"):
        block_id = int(name[len("down_blocks.")])
        hidden_size = unet.config.block_out_channels[block_id]
    layer_name = name.split(".processor")[0]
    weights = {
        "to_k_custom_diffusion.weight": st[layer_name + ".to_k.weight"],
        "to_v_custom_diffusion.weight": st[layer_name + ".to_v.weight"],
    }
    if train_q_out:
        weights["to_q_custom_diffusion.weight"] = st[layer_name + ".to_q.weight"]
        weights["to_out_custom_diffusion.0.weight"] = st[layer_name + ".to_out.0.weight"]
        weights["to_out_custom_diffusion.0.bias"] = st[layer_name + ".to_out.0.bias"]
    if cross_attention_dim is not None:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=train_kv,
            train_q_out=train_q_out,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        ).to(unet.device)
        custom_diffusion_attn_procs[name].load_state_dict(weights)
    else:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=False,
            train_q_out=False,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        )
del st
unet.set_attn_processor(custom_diffusion_attn_procs)
custom_diffusion_layers = AttnProcsLayers(unet.attn_processors)

optimizerは、クロスアテンション層のパラメータを更新するように初期化されます:

optimizer = optimizer_class(
    itertools.chain(text_encoder.get_input_embeddings().parameters(), custom_diffusion_layers.parameters())
    if args.modifier_token is not None
    else custom_diffusion_layers.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

学習ループでは、学習しようとしている概念の埋め込みのみを更新することが重要です。つまり、それ以外のすべてのトークン埋め込みの勾配をゼロに設定する必要があります。

if args.modifier_token is not None:
    if accelerator.num_processes > 1:
        grads_text_encoder = text_encoder.module.get_input_embeddings().weight.grad
    else:
        grads_text_encoder = text_encoder.get_input_embeddings().weight.grad
    index_grads_to_zero = torch.arange(len(tokenizer)) != modifier_token_id[0]
    for i in range(len(modifier_token_id[1:])):
        index_grads_to_zero = index_grads_to_zero & (
            torch.arange(len(tokenizer)) != modifier_token_id[i]
        )
    grads_text_encoder.data[index_grads_to_zero, :] = grads_text_encoder.data[
        index_grads_to_zero, :
    ].fill_(0)

スクリプトを実行します

すべての変更が完了したか、デフォルトの設定で問題ない場合は、トレーニングスクリプトの起動準備が整いました! 🚀

このガイドでは、これらの例として用意された猫の画像をダウンロードして使用します。必要に応じて、独自のデータセットを作成して使用することもできます（学習用データセットの作成ガイドを参照）。

環境変数 MODEL_NAME に Hub 上のモデル ID またはローカルモデルへのパスを、INSTANCE_DIR に先ほど猫の画像をダウンロードしたパスを、OUTPUT_DIR にモデルを保存したい場所を設定します。新たに学習した埋め込みを紐付けるための特別な単語として <new1> を使用します。スクリプトは、モデルのチェックポイントと pytorch_custom_diffusion_weights.bin ファイルを作成してリポジトリに保存します。

Weights and Biases でトレーニングの進捗を監視するには、トレーニングコマンドに --report_to=wandb パラメータを追加し、--validation_prompt で検証プロンプトを指定します。これにより、デバッグや中間結果の保存が可能になります。

人間の顔で学習を行う場合、Custom Diffusionチームは以下のパラメータが効果的であることを確認しています。

--learning_rate=5e-6

--max_train_steps は 1000 から 2000 の間の任意の値に設定できます

--freeze_model=crossattn

学習には少なくとも15〜20枚の画像を使用してください

単一の概念

複数の概念

トレーニングが完了したら、新しいCustom Diffusionモデルを推論に使用できます。

単一の概念

複数の概念

次のステップ

Custom Diffusionでモデルの学習を完了したことをお祝いします！🎉 詳細については、以下を参照してください：

Read the Multi-Concept Customization of Text-to-Image Diffusion blog post to learn more details about the experimental results from the Custom Diffusion team.

Update on GitHub

Diffusers

Custom Diffusion

スクリプトのパラメータ

事前分布の保存損失

正則化

学習スクリプト

スクリプトを実行します

次のステップ