テキストから画像へ

テキストから画像を生成するスクリプトは実験的なものであり、過学習やカタストロフィック・フォーゲッティングなどの問題が発生しやすいです。データセットで最適な結果を得るために、さまざまなハイパーパラメータの探索をお試しください。

Stable Diffusion などのテキストから画像を生成するモデルは、テキストプロンプトを条件として画像を生成するように学習されています。

モデルの学習はハードウェアに大きな負荷をかけますが、gradient_checkpointingとmixed_precisionを有効にすれば、24GBのGPU1枚でモデルを学習させることが可能です。より大きなバッチサイズで学習を行いたい場合や、学習速度を向上させたい場合は、30GB以上のメモリを搭載したGPUを使用することをお勧めします。xFormersでメモリ効率の高いアテンションを有効にすることで、メモリ使用量を削減できます。

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

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

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

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

cd examples/text_to_image
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_text_to_image.py \
  --mixed_precision="fp16"

基本的かつ重要なパラメータには、以下が含まれます。

--pretrained_model_name_or_path: Hub 上のモデル名、または事前学習済みモデルへのローカルパス
--dataset_name: Hub上のデータセット名、または学習に使用するデータセットのローカルパス
--image_column: 学習に使用するデータセット内の画像カラムの名前
--caption_column: 学習に使用するデータセット内のテキスト列の名前
--output_dir: 学習済みモデルを保存する場所
--push_to_hub: 学習したモデルを Hub にプッシュするかどうか
--checkpointing_steps: モデルの学習中にチェックポイントを保存する頻度です。学習が何らかの理由で中断された場合、学習コマンドに--resume_from_checkpointを追加することで、そのチェックポイントから学習を再開できます。

Min-SNR重み付け

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

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

accelerate launch train_text_to_image.py \
  --snr_gamma=5.0

このWeights and Biasesレポートでは、異なるsnr_gamma値における損失表面を比較できます。小規模なデータセットでは、大規模なデータセットと比較して、Min-SNRの影響がそれほど明白ではない場合があります。

学習スクリプト

データセットの前処理コードと学習ループは、main() 関数内にあります。学習スクリプトを調整する必要がある場合は、ここで変更を加える必要があります。

train_text_to_image スクリプトは、スケジューラーとトークナイザーを読み込むことから始まります。必要に応じて、ここで別のスケジューラーを使用することもできます。

noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
tokenizer = CLIPTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
)

次に、スクリプトはUNetモデルを読み込みます:

load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
model.register_to_config(**load_model.config)

model.load_state_dict(load_model.state_dict())

次に、データセットのテキスト列と画像列の前処理が必要です。tokenize_captions 関数は入力のトークン化を処理し、train_transforms 関数は画像に適用する変換の種類を指定します。これら2つの関数は preprocess_train にまとめられています。

def preprocess_train(examples):
    images = [image.convert("RGB") for image in examples[image_column]]
    examples["pixel_values"] = [train_transforms(image) for image in images]
    examples["input_ids"] = tokenize_captions(examples)
    return examples

最後に、training loopがそれ以外のすべての処理を担当します。画像を潜在空間にエンコードし、潜在変数にノイズを追加し、条件付けするためのテキスト埋め込みを計算し、モデルのパラメータを更新し、モデルを保存してHubにプッシュします。training loopの動作について詳しく知りたい場合は、デノイジングプロセスの基本的なパターンを解説しているUnderstanding pipelines, models and schedulersチュートリアルをご覧ください。

スクリプトを実行します

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

Naruto BLIP captions データセットで学習し、独自のナルトキャラクターを生成しましょう。環境変数 MODEL_NAME と dataset_name に、モデルとデータセット（Hub 由来またはローカルパス）を設定します。複数の GPU で学習する場合は、accelerate launch コマンドに --multi_gpu パラメータを追加してください。

ローカルデータセットで学習するには、TRAIN_DIR と OUTPUT_DIR 環境変数を、それぞれデータセットのパスとモデルの保存先パスに設定します。

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export dataset_name="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$dataset_name \
  --use_ema \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --enable_xformers_memory_efficient_attention \
  --lr_scheduler="constant" --lr_warmup_steps=0 \
  --output_dir="sd-naruto-model" \
  --push_to_hub

トレーニングが完了すると、新しくトレーニングしたモデルを推論に使用できます。

from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained("path/to/saved_model", dtype=torch.float16, use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"

image = pipeline(prompt="yoda").images[0]
image.save("yoda-naruto.png")

次のステップ

ご自身のテキストから画像を生成するモデルの学習が完了したことをお祝いします！新しいモデルの使用方法について詳しく知るために、以下のガイドが役立つ可能性があります：

LoRAでモデルを学習させた場合、推論のためにLoRAの重みをロードする方法を学びましょう。
guidance scale などの特定のパラメータや、prompt weighting などのテクニックが、Text-to-imageタスクガイドで説明されている推論の制御にどのように役立つのかを詳しく学びましょう。

Update on GitHub

Diffusers

テキストから画像へ

スクリプトのパラメータ

Min-SNR重み付け

学習スクリプト

スクリプトを実行します

次のステップ