DreamBooth

DreamBoothは、対象やスタイルの画像をわずか数枚だけ用いて学習し、拡散モデル全体を更新する学習手法です。プロンプト内の特別な単語を例示画像と関連付けることで機能します。

推論用に学習済みチェックポイントを読み込む場合は、推論用のDreamBoothモデルの読み込みを参照してください。

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

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

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

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

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

cd examples/dreambooth
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()

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

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

スクリプトのパラメータ

DreamBoothは学習ハイパーパラメータに非常に敏感であり、オーバーフィッティングを起こしやすいです。適切なハイパーパラメータを選択するために、🧨 Diffusers を使用して Stable Diffusion を Dreambooth で学習するというブログ記事で、異なる被写体向けの推奨設定を確認してください。

トレーニングスクリプトには、トレーニング実行をカスタマイズするための多くのパラメータが用意されています。すべてのパラメータとその説明は、parse_args()関数に記載されています。これらのパラメータには、そのままでも概ね良好な結果が得られるデフォルト値が設定されていますが、必要に応じてトレーニングコマンド内で独自の値を設定することもできます。

例えば、bf16形式で学習を行うには:

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

知っておくべき、かつ指定すべき基本的で重要なパラメータは以下の通りです。

--pretrained_model_name_or_path: Hub 上のモデル名、または事前学習済みモデルへのローカルパス
--instance_data_dir: 学習データセット（例画像）を含むフォルダへのパス
--instance_prompt: 例示画像に使用される特別な単語を含むテキストプロンプト
--train_text_encoder: テキストエンコーダーも同時に学習するかどうか
--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_dreambooth.py \
  --snr_gamma=5.0

事前分布の保存損失

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

--with_prior_preservation: 事前分布保持損失を使用するかどうか
--prior_loss_weight: モデルに対する事前分布保持損失の影響度を制御します
--class_data_dir: 生成されたクラスサンプル画像を含むフォルダへのパス
--class_prompt: 生成されるサンプル画像のクラスを説明するテキストプロンプト

accelerate launch train_dreambooth.py \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --class_data_dir="path/to/class/images" \
  --class_prompt="text prompt describing class"

テキストエンコーダーの学習

生成される出力の品質を向上させるには、UNetに加えてテキストエンコーダーも学習させることができます。これには追加のメモリが必要であり、少なくとも24GBのvRAMを備えたGPUが求められます。必要なハードウェアを所有している場合、テキストエンコーダーの学習により、特に顔の画像生成においてより良い結果が得られます。このオプションを有効にするには、次の手順を実行します。

accelerate launch train_dreambooth.py \
  --train_text_encoder

学習スクリプト

DreamBoothには独自のデータセットクラスが用意されています:

DreamBoothDataset: 学習用の画像とクラス画像の前処理、およびプロンプトのトークン化を行います
PromptDataset: クラス画像を生成するためのプロンプト埋め込みを生成します

prior preservation lossを有効にした場合、クラス画像はここで生成されます:

sample_dataset = PromptDataset(args.class_prompt, num_new_images)
sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=args.sample_batch_size)

sample_dataloader = accelerator.prepare(sample_dataloader)
pipeline.to(accelerator.device)

for example in tqdm(
    sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process
):
    images = pipeline(example["prompt"]).images

次に、main() 関数です。この関数は、学習用のデータセットのセットアップと、学習ループ自体を処理します。スクリプトは、tokenizer、schedulerとモデルを読み込みます。

# Load the tokenizer
if args.tokenizer_name:
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, revision=args.revision, use_fast=False)
elif args.pretrained_model_name_or_path:
    tokenizer = AutoTokenizer.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="tokenizer",
        revision=args.revision,
        use_fast=False,
    )

# Load scheduler and models
noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
text_encoder = text_encoder_cls.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
)

if model_has_vae(args):
    vae = AutoencoderKL.from_pretrained(
        args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision
    )
else:
    vae = None

unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
)

次に、DreamBoothDatasetから学習データセットとDataLoaderを作成します。

train_dataset = DreamBoothDataset(
    instance_data_root=args.instance_data_dir,
    instance_prompt=args.instance_prompt,
    class_data_root=args.class_data_dir if args.with_prior_preservation else None,
    class_prompt=args.class_prompt,
    class_num=args.num_class_images,
    tokenizer=tokenizer,
    size=args.resolution,
    center_crop=args.center_crop,
    encoder_hidden_states=pre_computed_encoder_hidden_states,
    class_prompt_encoder_hidden_states=pre_computed_class_prompt_encoder_hidden_states,
    tokenizer_max_length=args.tokenizer_max_length,
)

train_dataloader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=args.train_batch_size,
    shuffle=True,
    collate_fn=lambda examples: collate_fn(examples, args.with_prior_preservation),
    num_workers=args.dataloader_num_workers,
)

最後に、training loopが、画像を潜在空間に変換すること、入力にノイズを追加すること、ノイズ残差を予測すること、損失を計算することなどの残りのステップを担当します。

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

スクリプトを実行します

トレーニングスクリプトの起動準備が整いました! 🚀

このガイドでは、犬の画像をダウンロードしてディレクトリに保存します。ただし、必要に応じて独自のデータセットを作成して使用することもできます（学習用データセットの作成ガイドを参照）。

from huggingface_hub import snapshot_download

local_dir = "./dog"
snapshot_download(
    "diffusers/dog-example",
    local_dir=local_dir,
    repo_type="dataset",
    ignore_patterns=".gitattributes",
)

環境変数 MODEL_NAME に Hub 上のモデル ID またはローカルモデルへのパスを、INSTANCE_DIR に先ほど犬の画像をダウンロードしたパスを、OUTPUT_DIR にモデルを保存したい場所を設定します。学習を紐付けるための特別な単語として sks を使用します。

トレーニングの進捗を追跡したい場合は、トレーニング中に定期的に生成された画像を保存できます。トレーニングコマンドに次のパラメータを追加してください。

--validation_prompt="a photo of a sks dog"
--num_validation_images=4
--validation_steps=100

スクリプトを実行する前に、もう1点だけ確認事項があります。お使いのGPUによっては、DreamBoothの学習のために特定の最適化を有効にする必要がある場合があります。

16GB

12GB

8GB

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export INSTANCE_DIR="./dog"
export OUTPUT_DIR="path_to_saved_model"

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of sks dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=400 \
  --push_to_hub

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

学習が完了する前に推論でモデルを試してみたいですか? 🤭 🤗 Accelerate の最新バージョンがインストールされていることを確認してください。

from diffusers import DiffusionPipeline, UNet2DConditionModel
from transformers import CLIPTextModel
import torch

unet = UNet2DConditionModel.from_pretrained("path/to/model/checkpoint-100/unet")

# if you have trained with `--args.train_text_encoder` make sure to also load the text encoder
text_encoder = CLIPTextModel.from_pretrained("path/to/model/checkpoint-100/checkpoint-100/text_encoder")

pipeline = DiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, text_encoder=text_encoder, dtype=torch.float16,
).to("cuda")

image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0]
image.save("dog-bucket.png")

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("path_to_saved_model", dtype=torch.float16, use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"
image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0]
image.save("dog-bucket.png")

LoRA

LoRAは、学習可能なパラメータ数を大幅に削減するための学習手法です。その結果、学習が高速化され、生成される重みは非常に小さい（約100MB）ため、保存しやすくなります。LoRAを用いた学習には、train_dreambooth_lora.pyスクリプトを使用してください。

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

Stable Diffusion XL

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

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

DeepFloyd IF

DeepFloyd IF は、3段階の段階的ピクセル拡散モデルです。第1段階でベース画像を生成し、第2段階と第3段階でベース画像を段階的にアップスケールして、高解像度の1024x1024画像に仕上げます。train_dreambooth_lora.py または train_dreambooth.py スクリプトを使用して、LoRAまたはフルモデルで DeepFloyd IF モデルを学習できます。

DeepFloyd IF は予測分散を使用していますが、Diffusers の学習スクリプトは予測誤差を使用しているため、学習済みの DeepFloyd IF モデルは固定分散スケジュールに切り替えられます。学習スクリプトは、完全に学習されたモデルのスケジューラ設定を自動的に更新します。ただし、保存された LoRA 重みをロードする際は、パイプラインのスケジューラ設定も手動で更新する必要があります。

from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", use_safetensors=True)

pipe.load_lora_weights("<lora weights path>")

# Update scheduler config to fixed variance schedule
pipe.scheduler = pipe.scheduler.__class__.from_config(pipe.scheduler.config, variance_type="fixed_small")

ステージ2のモデルは、アップスケールするために追加の検証画像が必要です。このために、トレーニング画像の縮小版をダウンロードして使用できます。

from huggingface_hub import snapshot_download

local_dir = "./dog_downsized"
snapshot_download(
    "diffusers/dog-example-downsized",
    local_dir=local_dir,
    repo_type="dataset",
    ignore_patterns=".gitattributes",
)

以下のコードサンプルは、DreamBoothとLoRAの組み合わせでDeepFloyd IFモデルを学習する方法を簡潔に示しています。特に注意すべき重要なパラメータは以下の通りです。

--resolution=64、DeepFloyd IF はピクセル拡散モデルであり、非圧縮ピクセルを扱うため、入力画像はより小さいサイズである必要があります
--pre_compute_text_embeddings、T5Model は大量のメモリを消費する可能性があるため、メモリを節約するためにテキスト埋め込みを事前に計算します
--tokenizer_max_length=77 を指定すると、T5 をテキストエンコーダーとして使用する場合に、より長いデフォルトのテキスト長を使用できますが、デフォルトのモデルエンコーディング手順では、より短いテキスト長が使用されます。
--text_encoder_use_attention_mask を使用して、アテンションマスクをテキストエンコーダーに渡します。

Stage 1 LoRA DreamBooth

ステージ2 LoRA DreamBooth

ステージ1 DreamBooth

ステージ2 DreamBooth

LoRAとDreamBoothを用いたDeepFloyd IFのトレーニングステージ1には、約28GBのメモリが必要です。

export MODEL_NAME="DeepFloyd/IF-I-XL-v1.0"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="dreambooth_dog_lora"

accelerate launch train_dreambooth_lora.py \
  --report_to wandb \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a sks dog" \
  --resolution=64 \
  --train_batch_size=4 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --scale_lr \
  --max_train_steps=1200 \
  --validation_prompt="a sks dog" \
  --validation_epochs=25 \
  --checkpointing_steps=100 \
  --pre_compute_text_embeddings \
  --tokenizer_max_length=77 \
  --text_encoder_use_attention_mask

トレーニングのコツ

DeepFloyd IFモデルの学習は難しい場合がありますが、以下に有用だと感じたヒントをいくつかご紹介します。

モデルの解像度が低いため、細部の表現はどのみち困難であるため、ステージ1モデルの学習にはLoRAで十分です。
一般的かつシンプルなオブジェクトの場合、アップスケーラーのファインチューニングは必ずしも必要ありません。アップスケーラーに渡すプロンプトを調整し、インスタンスプロンプトから新しいトークンを削除するようにしてください。例えば、ステージ1のプロンプトが「a sks dog」である場合、ステージ2のプロンプトは「a dog」にする必要があります。
顔などの細部をより鮮明にするためには、ステージ2のアップスケーラーを完全に学習させる方が、LoRAを使用してステージ2のモデルを学習させるよりも効果的です。また、より低い学習率とより大きなバッチサイズを使用することも有効です。
ステージ2のモデルを学習させるには、より低い学習率を使用する必要があります。
DDPMSchedulerは、学習スクリプトで使用されているDPMSolverよりも優れた動作を示します。

次のステップ

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

推論用にDreamBoothのチェックポイントをロードする方法を学びましょう。
LoRAで学習を行った場合は、推論のためのLoRAを参照してください。

Update on GitHub