ディフューザを使用してControlNetをトレーニングしてください

Please train ControlNet using a diffuser.

イントロダクション

ControlNetは、追加の条件を付加することで拡散モデルを細かく制御することができるニューラルネットワーク構造です。この技術は、「Adding Conditional Control to Text-to-Image Diffusion Models」という論文で登場し、すぐにオープンソースの拡散コミュニティで広まりました。著者はStable Diffusion v1-5を制御するための8つの異なる条件をリリースしました。これには、ポーズ推定、深度マップ、キャニーエッジ、スケッチなどが含まれます。

このブログ投稿では、3Dシンセティックフェイスに基づいた顔のポーズモデルであるUncanny Facesモデルのトレーニング手順を詳細に説明します（実際にはUncanny Facesは予期しない結果であり、それがどのように実現されたかについては後ほどご紹介します）。

安定した拡散のためのControlNetのトレーニングの始め方

独自のControlNetをトレーニングするには、3つのステップが必要です：

• Intel CPU上での安定な拡散推論の高速化
• Substraを使用してプライバシーを保護するAIの作成
• トランスフォーマーによるグラフ分類

1. 条件の計画：ControlNetはStable Diffusionをさまざまなタスクに対応できる柔軟性があります。事前にトレーニングされたモデルはさまざまな条件を示しており、コミュニティはピクセル化されたカラーパレットに基づいた他の条件を作成しています。

2. データセットの構築：条件が決まったら、データセットの構築の時間です。そのためには、データセットをゼロから構築するか、既存のデータセットの一部を使用することができます。モデルをトレーニングするためには、データセットには3つの列が必要です：正解のimage、conditioning_image、およびprompt。

3. モデルのトレーニング：データセットの準備ができたら、モデルのトレーニングの時間です。これは、ディフューザーのトレーニングスクリプトのおかげで最も簡単な部分です。少なくとも8GBのVRAMを持つGPUが必要です。

1. 条件の計画

条件を計画するために、次の2つの質問を考えると役立ちます：

1. どのような条件を使用したいですか？
2. 既存のモデルで「通常の」画像を私の条件に変換できるものはありますか？

この例では、顔のランドマーク条件を使用することを考えました。その理由は次のとおりです：1. 一般的なランドマーク条件付きのControlNetはうまく機能します。2. 顔のランドマークは広く使用されており、通常の画像で顔のランドマークを計算するモデルが複数あります。3. Stable Diffusionを特定の顔のランドマークに従わせたり、自分自身の表情を模倣させるのは面白いかもしれません。

2. データセットの構築

では、顔のランドマークを持つStable Diffusion条件を作成することにしました。したがって、データセットを準備するためには以下が必要です：

• 正解のimage：この場合、顔の画像
• conditioning_image：この場合、顔のランドマークが可視化された画像
• caption：使用されている画像を説明するキャプション

このプロジェクトでは、MicrosoftのFaceSyntheticsデータセットを使用することにしました。これは、10万個の合成顔が含まれているデータセットです。他のリアルな顔の研究用データセットとしては、Celeb-A HQ、FFHQなどがありますが、このプロジェクトでは合成顔を選ぶことにしました。

FaceSyntheticsデータセットは素晴らしいスタート地点になると思われました。このデータセットには、顔の正解画像とiBUG 68-顔のランドマーク形式で注釈付けされた顔のセグメンテーション画像が含まれています。

完璧ですね？実際には、そうでもありません。”条件の計画”の手順での2番目の質問を思い出してください – つまり、通常の画像をこのデータセットの注釈付きランドマーク形式に変換できる既知のモデルが存在するはずです。しかし、実際にはこのデータセットの顔を注釈付きランドマーク形式に変換できる既知のモデルは存在しませんでした。

そこで、別の方法を選ぶことにしました：

• FaceSyntheticsデータセットの顔の正解imageを使用する
• 任意の顔の画像をiBUGの68フェイシャルランドマーク形式に変換できる既知のモデル（ここでは最先端のモデルSPIGAを使用）を使用する
• 顔のランドマークを見栄えの良いイラストのマスクに変換するカスタムコードを使用して、conditioning_imageとして使用する
• それをHugging Faceデータセットとして保存する

ここに、FaceSyntheticsデータセットから正解の画像をイラストのマスクに変換してHugging Faceデータセットとして保存するために使用されたコードがあります。

さて、データセットに正解のimageとconditioning_imageがあるため、あと1つのステップが足りません：各画像にキャプションが必要です。このステップは強くお勧めしますが、空のプロンプトで実験し、結果を報告することもできます。私たちはFaceSyntheticsデータセットにキャプションがなかったため、それをBLIPキャプションにかけました。すべての画像のキャプションに使用されたコードはこちらで確認できます

これで、最終的なデータセットに到着しました！ Face Synthetics SPIGA with captionsには、FaceSyntheticsデータセットの100Kの画像に対して正解の画像、セグメンテーション、キャプションが含まれています。モデルのトレーニングの準備が整いました！

3. モデルのトレーニング

準備ができたデータセットで、モデルのトレーニングの時間です！このプロセスの中で最も難しい部分であるはずでしたが、diffusersのトレーニングスクリプトを使用すると、実際には最も簡単な部分になりました。LambdaLabsで1.10ドル/時間で借りた単一のA100を使用しました。

私たちのトレーニングの経験

私たちはモデルを3エポックトレーニングしました（つまり、100Kの画像バッチがモデルに3回表示されました）し、バッチサイズは4でした（各ステップでモデルには4つの画像が表示されます）。これは過剰であり、過学習してしまいました（つまり、リアルな顔からやや外れた概念を忘れてしまい、たとえばプロンプトの「シュレック」や「猫」ではシュレックや猫ではなく、むしろ人物を生成するようになり、スタイルも無視し始めました）。

たった1エポック（つまり、モデルが100Kの画像を「見た」後）で、ポーズの追従が始まり、過学習しなくなりました。うまくいったのですが…、私たちはFace Syntheticsデータセットを使用したため、モデルは現実的な顔ではなく奇妙な3Dの顔を学びました。これは、合成の顔のデータセットを使用したためであり、楽しい/ミームの目的に使用できます。ここにuncannyfaces_25Kモデルがあります。

この対話型テーブルでは、下のダイヤルでモデルがどれくらいのトレーニングステップを経るかと、それがトレーニングプロセスにどのように影響するかを試すことができます。約15Kステップでポーズを学び始め、25Kステップで成熟しました。こちら

トレーニング方法は？

私たちがやるべきことは、依存関係をインストールするだけでした：

pip install git+https://github.com/huggingface/diffusers.git transformers accelerate xformers==0.0.16 wandb
huggingface-cli login
wandb login 

そして、train_controlnet.pyコードを実行するだけです

!accelerate launch train_controlnet.py \
 --pretrained_model_name_or_path="stabilityai/stable-diffusion-2-1-base" \
 --output_dir="model_out" \
 --dataset_name=multimodalart/facesyntheticsspigacaptioned \
 --conditioning_image_column=spiga_seg \
 --image_column=image \
 --caption_column=image_caption \
 --resolution=512 \
 --learning_rate=1e-5 \
 --validation_image "./face_landmarks1.jpeg" "./face_landmarks2.jpeg" "./face_landmarks3.jpeg" \
 --validation_prompt "High-quality close-up dslr photo of man wearing a hat with trees in the background" "Girl smiling, professional dslr photograph, dark background, studio lights, high quality" "Portrait of a clown face, oil on canvas, bittersweet expression" \
 --train_batch_size=4 \
 --num_train_epochs=3 \
 --tracker_project_name="controlnet" \
 --enable_xformers_memory_efficient_attention \
 --checkpointing_steps=5000 \
 --validation_steps=5000 \
 --report_to wandb \
 --push_to_hub

いくつかの設定を分解してみましょう。また、トレーニングに8GBのVRAMまで下げるための最適化のヒントも紹介します。

• pretrained_model_name_or_path：使用したい安定したディフュージョンベースモデル（ここではv2-1を選択しました。顔のレンダリングがより良くなります）
• output_dir：モデルを保存したいディレクトリ
• dataset_name：トレーニングに使用されるデータセット。私たちの場合は、「Face Synthetics SPIGA with captions」です
• conditioning_image_column：データセット内の条件画像を含む列の名前（私たちの場合はspiga_seg）
• image_column：データセット内の正解画像を含む列の名前（私たちの場合はimage）
• caption_column：データセット内の画像のキャプションを含む列の名前（私たちの場合はimage_caption）
• resolution：条件画像と正解画像の解像度（私たちの場合は512x512）
• learning_rate：学習率。これらの例では1e-5がうまく機能することがわかりましたが、1e-4から2e-6までの異なる値を試すこともできます。
• validation_image：トレーニング中に覗いてみるためのものです！バリデーションイメージは、validation_stepsごとに実行され、トレーニングの進捗状況がわかります。任意の数の条件画像へのローカルパスをここに挿入してください
• validation_prompt：バリデーションイメージと一緒に実行されるプロンプト。モデルがうまくトレーニングされているかテストできるものであれば何でも構いません
• train_batch_size：GPUに合わせてトレーニングバッチのサイズを指定します。A100を使用しているため、4とすることができますが、VRAMが小さいGPUを使用している場合はこの値を1にすることをお勧めします。
• num_train_epochs：各エポックは、トレーニングセットの画像がモデルに「見える」回数を表します。私たちは3エポックで実験しましたが、最良の結果には1エポック以上少し余分な時間が必要で、3エポックだとモデルが過学習します。
• checkpointing_steps：一定のステップ（ここでは5000）ごとに中間チェックポイントを保存します。5000ステップごとに、中間チェックポイントが保存されました。
• validation_steps：一定のステップ（xステップ）ごとにvalidaton_promptとvalidation_imageが実行されます。
• report_to：トレーニングを報告する場所。ここではWeights and Biasesを使用し、この素晴らしいレポートを得ました。ただし、train_batch_sizeを4から1に減らすだけでは、小さなGPUにトレーニングが適合しない場合があります。以下は各GPU VRAMサイズに追加するいくつかの追加パラメータです：
• push_to_hub：トレーニング済みモデルをHugging Face Hubにプッシュするためのパラメータ。

16GBのVRAMを持つGPUに適合させる

pip install bitsandbytes

--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--gradient_checkpointing \
--use_8bit_adam

バッチサイズ1と4つの勾配蓄積ステップの組み合わせは、例で使用した元のバッチサイズ4と同等です。さらに、勾配チェックポイントと8ビットAdamを有効にして、メモリの節約を行いました。

12GBのVRAMを持つGPUに適合させる

--gradient_accumulation_steps=4 \
--gradient_checkpointing \
--use_8bit_adam
--set_grads_to_none

8GB VRAM GPUでのフィッティング

こちらのガイドに従ってください

4. 結論！

ControlNetのトレーニングはとても楽しかったです。実際の顔のポーズを追従できるモデルを成功裏にトレーニングしましたが、実際の3Dの顔ではなく奇妙な3Dの顔を生成するように学習しました。これはトレーニングに使用されたデータセットが持つ独自の魅力と特徴です。

Hugging Face Spaceをお試しください：

次のステップとして、実際の顔のデータセットを使用せずに現実的に見える顔を作成するためのアイデアの1つは、FaceSyntheticsデータセット全体をStable Diffusion Image2Imageを介して実行し、3Dのように見える顔を現実的に見える顔に変換し、別のControlNetをトレーニングすることです。

そして、お楽しみに、ControlNetトレーニングイベントを間もなく開催いたします！Hugging FaceのTwitterをフォローするか、Discordに参加して最新情報を入手してください。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!