Stable Diffusion XL

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

Stable Diffusion XL (SDXL)は、Stable Diffusionモデルのより大型で高性能なバージョンであり、より高解像度の画像を生成できます。

SDXLのUNetは3倍大きく、アーキテクチャに第2のテキストエンコーダーが追加されています。利用可能なハードウェアによっては、計算負荷が非常に高くなる場合があり、Tesla T4のようなコンシューマー向けGPUでは実行できない可能性があります。この大きなモデルをメモリに収め、トレーニングを高速化するために、gradient_checkpointing、mixed_precision、gradient_accumulation_stepsを有効にすることを試してください。xFormersでメモリ効率の高いアテンションを有効にし、bitsandbytesの8ビットオプティマイザーを使用することで、メモリ使用量をさらに削減できます。

This guide will explore the train_text_to_image_sdxl.py training script to help you become more familiar with it, and how you can adapt it for your own use-case.

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

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

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

cd examples/text_to_image
pip install -r requirements_sdxl.txt

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

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

accelerate config

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

accelerate config default

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

from accelerate.utils import write_basic_config

write_basic_config()

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

スクリプトのパラメータ

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

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

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

accelerate launch train_text_to_image_sdxl.py \
  --mixed_precision="bf16"

ほとんどのパラメータはText-to-imageトレーニングガイドのパラメータと同一であるため、本ガイドではSDXLのトレーニングに関連するパラメータに焦点を当てます。

--pretrained_vae_model_name_or_path: 事前学習済み VAE のパス。SDXL の VAE は数値的不安定性を生じることが知られているため、このパラメータによりより優れた VAE を指定できます。
--proportion_empty_prompts: 画像プロンプトを空文字列に置き換える割合
--timestep_bias_strategy: バイアスを適用するタイムステップの位置（初期か後期か）を指定します。これにより、モデルが低周波数または高周波数の詳細を学習するように促すことができます。
--timestep_bias_multiplier: timestep に適用するバイアスの重み
--timestep_bias_begin: バイアスを適用し始めるタイムステップ
--timestep_bias_end: バイアスを適用する終了タイムステップ
--timestep_bias_portion: バイアスを適用するタイムステップの割合

Min-SNR重み付け

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

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

accelerate launch train_text_to_image_sdxl.py \
  --snr_gamma=5.0

学習スクリプト

学習スクリプトは、Text-to-imageの学習ガイドと似ていますが、SDXLの学習に対応するように変更されています。このガイドでは、SDXLの学習スクリプト固有のコードに焦点を当てます。

プロンプトの埋め込みを計算するためのプロンプトのトークン化、およびVAEを用いた画像埋め込みの計算を行う関数を作成することから始まります。次に、適用するタイムステップ数とタイムステップのバイアス戦略に応じて、タイムステップの重みを生成する関数を作成します。

main() 関数内では、トークナイザーの読み込みに加えて、SDXL アーキテクチャがそれぞれ2つを使用するため、2つ目のトークナイザーとテキストエンコーダーを読み込みます。

tokenizer_one = AutoTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision, use_fast=False
)
tokenizer_two = AutoTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer_2", revision=args.revision, use_fast=False
)

text_encoder_cls_one = import_model_class_from_model_name_or_path(
    args.pretrained_model_name_or_path, args.revision
)
text_encoder_cls_two = import_model_class_from_model_name_or_path(
    args.pretrained_model_name_or_path, args.revision, subfolder="text_encoder_2"
)

プロンプトと画像の埋め込みは最初に計算され、メモリに保持されます。これは小規模なデータセットでは通常問題になりませんが、大規模なデータセットではメモリ問題を引き起こす可能性があります。そのような場合は、事前計算された埋め込みをディスクに個別に保存し、トレーニングプロセス中にメモリへ読み込むことをお勧めします（このトピックに関する詳細な議論については、このPRを参照してください）。

text_encoders = [text_encoder_one, text_encoder_two]
tokenizers = [tokenizer_one, tokenizer_two]
compute_embeddings_fn = functools.partial(
    encode_prompt,
    text_encoders=text_encoders,
    tokenizers=tokenizers,
    proportion_empty_prompts=args.proportion_empty_prompts,
    caption_column=args.caption_column,
)

train_dataset = train_dataset.map(compute_embeddings_fn, batched=True, new_fingerprint=new_fingerprint)
train_dataset = train_dataset.map(
    compute_vae_encodings_fn,
    batched=True,
    batch_size=args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps,
    new_fingerprint=new_fingerprint_for_vae,
)

埋め込みを計算した後、メモリを解放するために、テキストエンコーダー、VAE、トークナイザーを削除します:

del text_encoders, tokenizers, vae
gc.collect()
torch.cuda.empty_cache()

最後に、training loopがそれ以外の処理を担当します。タイムステップのバイアス戦略を適用することを選んだ場合、タイムステップの重みが計算され、ノイズとして加算されることを確認できます。

weights = generate_timestep_weights(args, noise_scheduler.config.num_train_timesteps).to(
        model_input.device
    )
    timesteps = torch.multinomial(weights, bsz, replacement=True).long()

noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps)

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

スクリプトを実行します

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

Naruto BLIP captions データセットで学習し、独自のナルトキャラクターを生成しましょう。モデルとデータセット（Hub 由来またはローカルパス）をそれぞれ環境変数 MODEL_NAME と DATASET_NAME に設定します。また、数値的不安定さを避けるため、SDXL VAE 以外の VAE（Hub 由来またはローカルパス）を VAE_NAME で指定する必要があります。

Weights & Biases でトレーニングの進捗を監視するには、トレーニングコマンドに --report_to=wandb パラメータを追加します。また、結果を追跡するために、トレーニングコマンドに --validation_prompt と --validation_epochs も追加する必要があります。これにより、モデルのデバッグや中間結果の確認に非常に役立ちます。

export MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0"
export VAE_NAME="madebyollin/sdxl-vae-fp16-fix"
export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch train_text_to_image_sdxl.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --pretrained_vae_model_name_or_path=$VAE_NAME \
  --dataset_name=$DATASET_NAME \
  --enable_xformers_memory_efficient_attention \
  --resolution=512 \
  --center_crop \
  --random_flip \
  --proportion_empty_prompts=0.2 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=10000 \
  --use_8bit_adam \
  --learning_rate=1e-06 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --mixed_precision="fp16" \
  --report_to="wandb" \
  --validation_prompt="a cute Sundar Pichai creature" \
  --validation_epochs 5 \
  --checkpointing_steps=5000 \
  --output_dir="sdxl-naruto-model" \
  --push_to_hub

トレーニングが完了したら、新しくトレーニングしたSDXLモデルを使って推論を実行できます！

PyTorch

PyTorch XLA

次のステップ

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

Stable Diffusion XLのガイドを読み、さまざまなタスク（テキストから画像への変換、画像から画像への変換、インペインティング）での使い方、リファイナーモデルの使い方、および異なるタイプのマイクロコンディショニングについて学びましょう。
DreamBoothとLoRAのトレーニングガイドを参照し、わずか数枚のサンプル画像でパーソナライズされたSDXLモデルをトレーニングする方法を学びましょう。これらの2つのトレーニング手法は組み合わせることも可能です。

Update on GitHub

Diffusers

Stable Diffusion XL

スクリプトのパラメータ

Min-SNR重み付け

学習スクリプト

スクリプトを実行します

次のステップ