# Custom Diffusion

[Custom Diffusion](https://huggingface.co/papers/2212.04488)は、画像生成モデルをパーソナライズするための学習手法です。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](https://github.com/huggingface/diffusers/blob/main/examples/custom_diffusion/train_custom_diffusion.py) スクリプトを詳しく見ていき、その仕組みをより深く理解したうえで、ご自身のユースケースに合わせてどのようにカスタマイズできるかを解説します。

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

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

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

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

> [!TIP]
> 🤗 Accelerate は、複数 GPU/TPU での学習や混合精度学習を支援するためのライブラリです。ハードウェアと環境に基づいて、学習設定を自動的に構成します。詳細については、🤗 Accelerate の [クイックツアー](https://huggingface.co/docs/accelerate/quicktour) をご覧ください。

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

```bash
accelerate config
```

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

```bash
accelerate config default
```

または、ノートブックのように対話型シェルをサポートしていない環境では、以下を使用できます:

```py
from accelerate.utils import write_basic_config

write_basic_config()
```

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

> [!TIP]
> 以下のセクションでは、スクリプトの修正方法を理解するために重要な部分に焦点を当てていますが、スクリプトのすべての側面を詳細に説明するものではありません。さらに詳しく知りたい場合は、[スクリプト](https://github.com/huggingface/diffusers/blob/main/examples/custom_diffusion/train_custom_diffusion.py)をぜひお読みいただき、ご質問やご懸念事項があればお知らせください。

## スクリプトのパラメータ

トレーニングスクリプトには、トレーニング実行をカスタマイズするためのすべてのパラメータが含まれています。これらは[`parse_args()`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/custom_diffusion/train_custom_diffusion.py#L319)関数内で定義されています。この関数にはデフォルト値が設定されていますが、必要に応じてトレーニングコマンド内で独自の値を設定することもできます。

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

```bash
accelerate launch train_custom_diffusion.py \
  --resolution=256
```

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

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

### 事前分布の保存損失

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

事前保持損失のパラメータの多くは、[DreamBooth](dreambooth#prior-preservation-loss)のトレーニングガイドで説明されています。

### 正則化

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

```bash
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`: 過学習を防ぐために、実際の画像の小さなセットを使用するかどうか

```bash
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 \
```

## 学習スクリプト

> [!TIP]
> Custom Diffusionの学習スクリプトには、[DreamBooth](dreambooth#training-script)スクリプトと類似したコードが多く含まれています。本ガイドでは、Custom Diffusionに関連するコードに焦点を当てて説明します。

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

- [`CustomDiffusionDataset`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/custom_diffusion/train_custom_diffusion.py#L165): 学習用に画像、クラス画像、プロンプトを前処理します
- [`PromptDataset`](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/custom_diffusion/train_custom_diffusion.py#L148): クラス画像の生成用にプロンプトを準備します

次に、`modifier_token` が [トークナイザーに追加](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/custom_diffusion/train_custom_diffusion.py#L811)され、トークン ID に変換されます。その後、新しい `modifier_token` を考慮して、トークン埋め込みのサイズが調整されます。続いて、`modifier_token` の埋め込みは、`initializer_token` の埋め込みで初期化されます。テキストエンコーダーのパラメータはすべて凍結されますが、モデルが概念に関連付けようとしているのがトークン埋め込みであるため、トークン埋め込みのみは凍結されません。

```py
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](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/custom_diffusion/train_custom_diffusion.py#L911C3-L911C3)をattentionレイヤーに追加する必要があります。これは、attention weightsの形状とサイズを正しく設定し、各UNetブロックに適切な数のattention processorsを設定するために非常に重要なステップです。

```py
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](https://github.com/huggingface/diffusers/blob/84cd9e8d01adb47f046b1ee449fc76a0c32dc4e2/examples/custom_diffusion/train_custom_diffusion.py#L982)は、クロスアテンション層のパラメータを更新するように初期化されます:

```py
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,
)
```

[学習ループ](https://github.com/huggingface/diffusers/blob/84cd9e8d01adb47f046b1ee449fc76a0c32dc4e2/examples/custom_diffusion/train_custom_diffusion.py#L1048)では、学習しようとしている概念の埋め込みのみを更新することが重要です。つまり、それ以外のすべてのトークン埋め込みの勾配をゼロに設定する必要があります。

```py
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)
```

## スクリプトを実行します

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

このガイドでは、これらの例として用意された[猫の画像](https://www.cs.cmu.edu/~custom-diffusion/assets/data.zip)をダウンロードして使用します。必要に応じて、独自のデータセットを作成して使用することもできます（[学習用データセットの作成](create_dataset)ガイドを参照）。

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

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

> [!TIP]
> 人間の顔で学習を行う場合、Custom Diffusionチームは以下のパラメータが効果的であることを確認しています。
>
> - `--learning_rate=5e-6` - `--max_train_steps` は 1000 から 2000 の間で設定できます - `--freeze_model=crossattn` - 学習には少なくとも 15〜20 枚の画像を使用してください

```bash
export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export OUTPUT_DIR="path-to-save-model"
export INSTANCE_DIR="./data/cat"

accelerate launch train_custom_diffusion.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --class_data_dir=./real_reg/samples_cat/ \
  --with_prior_preservation \
  --real_prior \
  --prior_loss_weight=1.0 \
  --class_prompt="cat" \
  --num_class_images=200 \
  --instance_prompt="photo of a <new1> cat"  \
  --resolution=512  \
  --train_batch_size=2  \
  --learning_rate=1e-5  \
  --lr_warmup_steps=0 \
  --max_train_steps=250 \
  --scale_lr \
  --hflip  \
  --modifier_token "<new1>" \
  --validation_prompt="<new1> cat sitting in a bucket" \
  --report_to="wandb" \
  --push_to_hub
```

[JSON](https://github.com/adobe-research/custom-diffusion/blob/main/assets/concept_list.json)ファイルに学習すべき各コンセプトの詳細を指定することで、Custom Diffusionは複数のコンセプトを学習することもできます。

正則化に使用する実際の画像を収集するために、clip-retrievalを実行します:

```bash
pip install clip-retrieval
python retrieve.py --class_prompt {} --class_data_dir {} --num_class_images 200
```

その後、スクリプトを起動できます。

```bash
export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export OUTPUT_DIR="path-to-save-model"

accelerate launch train_custom_diffusion.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --output_dir=$OUTPUT_DIR \
  --concepts_list=./concept_list.json \
  --with_prior_preservation \
  --real_prior \
  --prior_loss_weight=1.0 \
  --resolution=512  \
  --train_batch_size=2  \
  --learning_rate=1e-5  \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --num_class_images=200 \
  --scale_lr \
  --hflip  \
  --modifier_token "<new1>+<new2>" \
  --push_to_hub
```

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

```py
import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet.load_attn_procs("path-to-save-model", weight_name="pytorch_custom_diffusion_weights.bin")
pipeline.load_textual_inversion("path-to-save-model", weight_name="<new1>.bin")

image = pipeline(
    "<new1> cat sitting in a bucket",
    num_inference_steps=100,
    guidance_scale=6.0,
    eta=1.0,
).images[0]
image.save("cat.png")
```

```py
import torch
from huggingface_hub.repocard import RepoCard
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4", dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"
model_id = "sayakpaul/custom-diffusion-cat-wooden-pot"
pipeline.unet.load_attn_procs(model_id, weight_name="pytorch_custom_diffusion_weights.bin")
pipeline.load_textual_inversion(model_id, weight_name="<new1>.bin")
pipeline.load_textual_inversion(model_id, weight_name="<new2>.bin")

image = pipeline(
    "the <new1> cat sculpture in the style of a <new2> wooden pot",
    num_inference_steps=100,
    guidance_scale=6.0,
    eta=1.0,
).images[0]
image.save("multi-subject.png")
```

## 次のステップ

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

- Read the [Multi-Concept Customization of Text-to-Image Diffusion](https://www.cs.cmu.edu/~custom-diffusion/) blog post to learn more details about the experimental results from the Custom Diffusion team.