本記事では、InstructPix2Pixと呼ばれる機械学習手法を用いて任意の画像をテキストで編集する方法をご紹介します。

出典: InstructPix2Pix: Learning to Follow Image Editing Instructions

InstructPix2Pix

概要

InstructPix2Pixは、テキストベースの指示により画像を編集する拡散モデルベースの画像編集技術です。

InstructPix2Pixではこのタスクを実現するため言語モデル(GPT-3)とText to Imageタスクを実現するStable Diffusionの事前学習済みモデルを用いて、画像編集のための大規模データセットを生成します。

条件付き拡散モデルInstructPix2Pixは生成したデータセットでトレーニングされ様々な入力画像とテキストによる指示に対して有効な編集結果を出力することを実現しています。

以下は、モデルのトレーニング概要を示しています。

1. (a)GPT-3を使用しInstructionとEdited Captionを生成
2. (b)StableDiffusionをPrompt to Promptを使用しInput/Edited Captionと画像のペアを生成
3. (c)上記により450,000をトレーニングデータセットを作成
4. (d)生成したデータセットでInstructPix2Pixをトレーニング

出典: InstructPix2Pix: Learning to Follow Image Editing Instructions

詳細はこちらの論文をご参照ください。

本記事では上記手法を用いて、任意の画像をテキストで指示し編集していきます。

デモ(Colaboratory)

それでは、実際に動かしながら画像編集を行います。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

また、下記から直接Google Colaboratoryで開くこともできます。

なお、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍やオンライン講座などがおすすめです。

環境セットアップ

それではセットアップしていきます。 Colaboratoryを開いたら下記を設定しGPUを使用するようにしてください。

はじめにライブラリをインストールします。

!pip install git+https://github.com/huggingface/diffusers.git@69c76173faaa3831cc4bc6f19b60b4ac8e9e4473
!pip install transformers==4.25.1 accelerate==0.15.0 safetensors==0.2.8

最後にライブラリをインポートします。

import PIL
import requests
import torch
from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler

device = 'cuda' if torch.cuda.is_available() else "cpu"
print("using device is", device)

以上で環境セットアップは完了です。

学習済みモデルのセットアップ

続いて、HuggingFaceから学習済みモデルをロードします。

# 学習済みモデルをHuggingFaceからロード
model_id = "timbrooks/instruct-pix2pix"
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="fp16", safety_checker=None)

# GPUメモリに配置
pipe.to("cuda")
pipe.enable_attention_slicing()

テスト画像のセットアップ

続いて、編集する任意の画像をWeb上から取得します。

!wget -c https://www.pakutaso.com/shared/img/thumb/yuseiookawa1971918_TP_V4.jpg \
      -O test01.jpg

取得した画像を表示します。

image = PIL.Image.open("./test01.jpg")
image = image.convert("RGB")
image

Pix2Pix

それでは、テキストを入力し、画像編集を実行します。

以下では、turn him into a man with glassesを入力し、眼鏡をかけるよう指示しています。

prompt = "turn him into a man with glasses"
pipe(prompt, image=image, num_inference_steps=100, image_guidance_scale=0.11).images[0]

出力結果は以下の通りです。画像中の男性が眼鏡をかけていることが確認できます。

まとめ

本記事では、InstructPix2Pixで任意画像の編集を行いました。

当初はText to ImageタスクのみであったStable Diffusionですが、Pix2Pixなど様々なタスクへと広がりを見せています。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
もう少し学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。

また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1.  論文 - InstructPix2Pix: Learning to Follow Image Editing Instructions

2. GitHub - instruct-pix2pix