ControlNet

ControlNetモデルは、別の事前学習済みモデルの上に訓練されたアダプターです。追加の入力画像でモデルを条件付けすることで、画像生成に対する制御度を高めることができます。入力画像には、cannyエッジ、深度マップ、人間のポーズなど、さまざまなものを使用できます。

vRAM が限られた GPU で学習を行う場合は、学習コマンドで gradient_checkpointing、gradient_accumulation_steps、mixed_precision パラメータを有効にすることを試してください。また、xFormers を使用してメモリ効率の高いアテンションを実装することで、メモリ使用量を削減することもできます。

このガイドでは、train_controlnet.py トレーニングスクリプトを詳しく見ていき、その内容に慣れるとともに、ご自身のユースケースに合わせてどのように適応させられるかを解説します。

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

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

次に、トレーニングスクリプトが含まれるサンプルフォルダに移動し、使用するスクリプトに必要な依存関係をインストールします:

cd examples/controlnet
pip install -r requirements.txt

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

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

accelerate config

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

accelerate config default

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

from accelerate.utils import write_basic_config

write_basic_config()

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

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

スクリプトのパラメータ

トレーニングスクリプトには、トレーニング実行をカスタマイズするための多くのパラメータが用意されています。すべてのパラメータとその説明は、parse_args()関数に記載されています。この関数は、トレーニングのバッチサイズや学習率などの各パラメータにデフォルト値を設定しますが、必要に応じてトレーニングコマンドで独自の値を設定することもできます。

例えば、fp16 形式の混合精度でトレーニングを高速化するには、トレーニングコマンドに --mixed_precision パラメータを追加します。

accelerate launch train_controlnet.py \
  --mixed_precision="fp16"

基本的かつ重要なパラメータの多くは、Text-to-imageのトレーニングガイドで説明されているため、本ガイドではControlNetに関連するパラメータに焦点を当てます。

--max_train_samples: 学習サンプルの数。より高速な学習のためにこれを下げることができますが、非常に大きなデータセットをストリーミングで扱う場合は、学習コマンドにこのパラメータと--streamingパラメータを含める必要があります。
--gradient_accumulation_steps: バックワードパスの前に蓄積する更新ステップ数。これにより、GPUメモリが通常処理できる以上の大きなバッチサイズで学習できます。

Min-SNR重み付け

Min-SNR重み付け戦略は、損失を再バランス化して収束を高速化することで、トレーニングを支援できます。トレーニングスクリプトはepsilon（ノイズ）またはv_predictionの予測をサポートしていますが、Min-SNRはどちらの予測タイプにも対応しています。この重み付け戦略はPyTorchでのみサポートされています。

--snr_gamma パラメータを追加し、推奨値である 5.0 に設定します:

accelerate launch train_controlnet.py \
  --snr_gamma=5.0

学習スクリプト

スクリプトのパラメータと同様に、Text-to-imageトレーニングガイドでは、トレーニングスクリプトの一般的な解説が提供されています。本ガイドでは、ControlNetスクリプトの関連する部分に焦点を当てて解説します。

学習スクリプトには、画像変換とキャプションのトークン化によってデータセットを前処理するための make_train_dataset 関数が含まれています。通常のキャプションのトークン化や画像変換に加えて、スクリプトには条件付け画像用の変換も含まれていることがわかります。

TPU上でデータセットをストリーミングしている場合、画像に最適化されていない🤗 Datasetsライブラリがボトルネックとなり、パフォーマンスが低下する可能性があります。最大のスループットを確保するために、WebDataset、TorchData、TensorFlow Datasetsなどの他のデータセット形式を検討することをお勧めします。

conditioning_image_transforms = transforms.Compose(
    [
        transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
        transforms.CenterCrop(args.resolution),
        transforms.ToTensor(),
    ]
)

main() 関数内には、tokenizer、text encoder、scheduler、モデルを読み込むコードが含まれています。また、ControlNet モデルを既存の重みから読み込むか、UNet からランダムに初期化して読み込むかの処理もここで行われます。

if args.controlnet_model_name_or_path:
    logger.info("Loading existing controlnet weights")
    controlnet = ControlNetModel.from_pretrained(args.controlnet_model_name_or_path)
else:
    logger.info("Initializing controlnet weights from unet")
    controlnet = ControlNetModel.from_unet(unet)

optimizerは、ControlNetのパラメータを更新するように設定されています:

params_to_optimize = controlnet.parameters()
optimizer = optimizer_class(
    params_to_optimize,
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

最後に、学習ループでは、条件付けテキストの埋め込みと画像が、ControlNet モデルのダウンブロックとミッドブロックに渡されます。

encoder_hidden_states = text_encoder(batch["input_ids"])[0]
controlnet_image = batch["conditioning_pixel_values"].to(dtype=weight_dtype)

down_block_res_samples, mid_block_res_sample = controlnet(
    noisy_latents,
    timesteps,
    encoder_hidden_states=encoder_hidden_states,
    controlnet_cond=controlnet_image,
    return_dict=False,
)

トレーニングループの仕組みについてさらに詳しく知りたい場合は、デノイジングプロセスの基本的なパターンを解説しているパイプライン、モデル、スケジューラーの理解チュートリアルをご覧ください。

スクリプトを実行します

これで学習スクリプトの起動準備が整いました! 🚀

このガイドではfusing/fill50kデータセットを使用しますが、必要に応じて独自のデータセットを作成して使用することもできます（学習用データセットの作成ガイドを参照してください）。

環境変数 MODEL_NAME に Hub 上のモデル ID またはローカルモデルへのパスを、OUTPUT_DIR にモデルを保存する先を設定してください。

トレーニングの条件付けに使用する次の画像をダウンロードします:

wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_1.png
wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_2.png

スクリプトを実行する前に、もう1点確認事項があります。お使いのGPUによっては、ControlNetの学習のために特定の最適化を有効にする必要がある場合があります。このスクリプトのデフォルト設定では、約38GBのvRAMが必要です。複数のGPUで学習を行う場合は、accelerate launchコマンドに--multi_gpuパラメータを追加してください。

16GB

12GB

8GB

export MODEL_DIR="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="path/to/save/model"

accelerate launch train_controlnet.py \
 --pretrained_model_name_or_path=$MODEL_DIR \
 --output_dir=$OUTPUT_DIR \
 --dataset_name=fusing/fill50k \
 --resolution=512 \
 --learning_rate=1e-5 \
 --validation_image "./conditioning_image_1.png" "./conditioning_image_2.png" \
 --validation_prompt "red circle with blue background" "cyan circle with brown floral background" \
 --train_batch_size=1 \
 --gradient_accumulation_steps=4 \
 --push_to_hub

トレーニングが完了したら、新しく学習したモデルを推論に使用できます!

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers.utils import load_image
import torch

controlnet = ControlNetModel.from_pretrained("path/to/controlnet", dtype=torch.float16)
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
    "path/to/base/model", controlnet=controlnet, dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

control_image = load_image("./conditioning_image_1.png")
prompt = "pale golden rod circle with old lace background"

generator = torch.manual_seed(0)
image = pipeline(prompt, num_inference_steps=20, generator=generator, image=control_image).images[0]
image.save("./output.png")

Stable Diffusion XL

Stable Diffusion XL (SDXL) は高解像度の画像を生成する強力なテキストから画像へのモデルであり、そのアーキテクチャに第2のテキストエンコーダーを追加しています。SDXLモデル用のControlNetアダプターを学習するには、train_controlnet_sdxl.py スクリプトを使用します。

SDXLの学習スクリプトについては、SDXLの学習ガイドで詳しく説明しています。

次のステップ

ご自身のControlNetの学習が完了したことをお祝いします。新しいモデルの使用方法について詳しく知るために、以下のガイドが役立つ可能性があります：

さまざまなタスクで推論を行う際にControlNetを使用する方法を学びましょう。

Update on GitHub

Diffusers

ControlNet

スクリプトのパラメータ

Min-SNR重み付け

学習スクリプト

スクリプトを実行します

Stable Diffusion XL

次のステップ