Diffusers documentation

InstructPix2Pix

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

InstructPix2Pix

InstructPix2Pixは、人間の指示に基づいて画像を編集するように学習されたStable Diffusionモデルです。例えば、プロンプトに「雲を雨雲に変えて」と入力すると、モデルは入力画像をそれに従って編集します。このモデルは、テキストプロンプト(または編集指示)と入力画像の両方を条件として動作します。

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

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

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

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

cd examples/instruct_pix2pix
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()関数に記載されています。ほとんどのパラメータには、かなり良好に動作するデフォルト値が設定されていますが、必要に応じてトレーニングコマンドで独自の値を設定することもできます。

例えば、入力画像の解像度を上げるには:

accelerate launch train_instruct_pix2pix.py \
  --resolution=512 \

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

  • --original_image_column: 編集前の元の画像
  • --edited_image_column: 編集が適用された後の画像
  • --edit_prompt_column: 画像を編集するための指示
  • --conditioning_dropout_prob: 学習時に編集済み画像および編集プロンプトに対して適用されるドロップアウト確率。これにより、1つまたは両方の条件付け入力に対してclassifier-free guidance (CFG)を有効にできます

学習スクリプト

データセットの前処理コードと学習ループは、main() 関数内にあります。ここが、ご自身のユースケースに合わせて学習スクリプトを調整するための変更箇所です。

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

スクリプトは、InstructPix2Pixの追加の条件付け画像に対応するために、UNetの最初の畳み込み層の入力チャネル数を変更することから始まります:

in_channels = 8
out_channels = unet.conv_in.out_channels
unet.register_to_config(in_channels=in_channels)

with torch.no_grad():
    new_conv_in = nn.Conv2d(
        in_channels, out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding
    )
    new_conv_in.weight.zero_()
    new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight)
    unet.conv_in = new_conv_in

これらの UNet パラメータは、オプティマイザーによって更新されます:

optimizer = optimizer_cls(
    unet.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

次に、編集された画像と編集指示が前処理およびトークン化されます。元の画像と編集された画像の両方に同じ画像変換を適用することが重要です。

def preprocess_train(examples):
    preprocessed_images = preprocess_images(examples)

    original_images, edited_images = preprocessed_images.chunk(2)
    original_images = original_images.reshape(-1, 3, args.resolution, args.resolution)
    edited_images = edited_images.reshape(-1, 3, args.resolution, args.resolution)

    examples["original_pixel_values"] = original_images
    examples["edited_pixel_values"] = edited_images

    captions = list(examples[edit_prompt_column])
    examples["input_ids"] = tokenize_captions(captions)
    return examples

最後に、training loopでは、編集済み画像を潜在空間にエンコードすることから始めます。

latents = vae.encode(batch["edited_pixel_values"].to(weight_dtype)).latent_dist.sample()
latents = latents * vae.config.scaling_factor

その後、スクリプトは CFG をサポートするために、元の画像と編集指示の埋め込みに dropout を適用します。これにより、モデルは編集指示と元の画像が編集後の画像に与える影響を調整できるようになります。

encoder_hidden_states = text_encoder(batch["input_ids"])[0]
original_image_embeds = vae.encode(batch["original_pixel_values"].to(weight_dtype)).latent_dist.mode()

if args.conditioning_dropout_prob is not None:
    random_p = torch.rand(bsz, device=latents.device, generator=generator)
    prompt_mask = random_p < 2 * args.conditioning_dropout_prob
    prompt_mask = prompt_mask.reshape(bsz, 1, 1)
    null_conditioning = text_encoder(tokenize_captions([""]).to(accelerator.device))[0]
    encoder_hidden_states = torch.where(prompt_mask, null_conditioning, encoder_hidden_states)

    image_mask_dtype = original_image_embeds.dtype
    image_mask = 1 - (
        (random_p >= args.conditioning_dropout_prob).to(image_mask_dtype)
        * (random_p < 3 * args.conditioning_dropout_prob).to(image_mask_dtype)
    )
    image_mask = image_mask.reshape(bsz, 1, 1, 1)
    original_image_embeds = image_mask * original_image_embeds

ほぼこれで完了です! ここに記載した違いを除き、スクリプトの残りはテキストから画像への学習スクリプトと非常に似ており、詳細についてはそちらを参照してください。学習ループの仕組みについてさらに詳しく知りたい場合は、デノイジングプロセスの基本的なパターンを解説しているパイプライン、モデル、スケジューラーの理解チュートリアルをご覧ください。

スクリプトを実行します

スクリプトの変更内容に満足した、またはデフォルトの設定で問題ない場合は、トレーニングスクリプトの起動準備が整いました! 🚀

このガイドでは、元のデータセットの小型版であるfusing/instructpix2pix-1000-samplesデータセットを使用します。必要に応じて、ご自身のデータセットを作成して使用することもできます(学習用データセットの作成ガイドを参照)。

MODEL_NAME 環境変数にモデル名(Hub 上のモデル ID またはローカルモデルへのパス)を設定し、DATASET_ID には Hub 上のデータセット名を設定します。このスクリプトは、リポジトリ内のサブフォルダーにすべてのコンポーネント(feature extractor、scheduler、text encoder、UNet など)を作成して保存します。

より良い結果を得るには、より大きなデータセットを用いた長時間のトレーニングを実行することをお勧めします。本トレーニングスクリプトは、小規模なデータセットでのみテストを行っています。


Weights and Biases でトレーニングの進捗を監視するには、トレーニングコマンドに --report_to=wandb パラメータを追加し、--val_image_url で検証用画像、--validation_prompt で検証用プロンプトを指定します。これにより、モデルのデバッグに非常に役立ちます。

複数の GPU で学習を行う場合は、accelerate launch コマンドに --multi_gpu パラメータを追加してください。

accelerate launch --mixed_precision="fp16" train_instruct_pix2pix.py \
    --pretrained_model_name_or_path=$MODEL_NAME \
    --dataset_name=$DATASET_ID \
    --enable_xformers_memory_efficient_attention \
    --resolution=256 \
    --random_flip \
    --train_batch_size=4 \
    --gradient_accumulation_steps=4 \
    --gradient_checkpointing \
    --max_train_steps=15000 \
    --checkpointing_steps=5000 \
    --checkpoints_total_limit=1 \
    --learning_rate=5e-05 \
    --max_grad_norm=1 \
    --lr_warmup_steps=0 \
    --conditioning_dropout_prob=0.05 \
    --mixed_precision=fp16 \
    --seed=42 \
    --push_to_hub

トレーニングが完了すると、新しいInstructPix2Pixを使って推論を実行できます。

import PIL
import requests
import torch
from diffusers import StableDiffusionInstructPix2PixPipeline
from diffusers.utils import load_image

pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained("your_cool_model", dtype=torch.float16).to("cuda")  # or "mps", "xpu", "cpu"
generator = torch.Generator("cuda").manual_seed(0)

image = load_image("https://huggingface.co/datasets/sayakpaul/sample-datasets/resolve/main/test_pix2pix_4.png")
prompt = "add some ducks to the lake"
num_inference_steps = 20
image_guidance_scale = 1.5
guidance_scale = 10

edited_image = pipeline(
   prompt,
   image=image,
   num_inference_steps=num_inference_steps,
   image_guidance_scale=image_guidance_scale,
   guidance_scale=guidance_scale,
   generator=generator,
).images[0]
edited_image.save("edited_image.png")

推論の速度と品質への影響を確認するために、num_inference_steps、image_guidance_scale、guidance_scale の異なる値を試すことをお勧めします。guidance scale パラメータは、元の画像や編集指示が編集後の画像にどの程度影響を与えるかを制御するため、特に大きな影響を与えます。

Stable Diffusion XL

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

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

次のステップ

ご自身のInstructPix2Pixモデルの学習が完了しました!🥳 モデルについてさらに詳しく知るために、以下の方法が役立つ場合があります:

Update on GitHub