Stable Diffusionによる画像修復（image inpainting）

この記事では、Stable Diffusionを用いて、画像の指定領域をテキストによって修復（inpainting）する方法を紹介します。

Stable Diffusionを用いたtext-to-image、image-to-imageの方法は以下の記事をご覧ください。

Stable Diffusionとは
デモ
まとめ
参考文献

Stable Diffusionとは

Stable Diffusionは、High-Resolution Image Synthesis with Latent Diffusion Modelsで提案されたLatent Diffusionという特殊なタイプのDiffusionモデル（拡散モデル）に基づく画像生成モデルです。

拡散モデルは、画像生成タスクにおいて非常に優れた精度を発揮することが知られていますが、一般的にメモリ消費量が膨大であるため、学習や推論が困難でした。

Latent Diffusionは、GANの一種と同様に低次元の潜在表現を活用することでこの点を克服し、メモリの消費を軽減することに成功したため、より汎用的なモデルとなっています。

デモ

それでは、Google Colaboratoryを使って実際に実行していきます。

なお、記事内で紹介したコードをすべて含むノートブックは、以下のリンクから直接参照することができます。

事前準備

まず、モデルが公開されているHugging Faceのアカウントを作成します。

Hugging Faceのトップページから右上の「Sign Up」に進みます。アドレスとパスワードを入力すると、登録したアドレスに認証メールが届くので、メール内のリンクをクリックして作成は完了です。

次に、Stable Diffusionのリポジトリを開き「Access repository」をクリックします。

最後に、学習済みモデルをダウンロードするためのアクセストークンを取得します。

トップページから先ほど作成したアカウントでログイン後、右上のアカウントのアイコンをクリックし、「Settings」、「Access Tokens」と進みます。画面中央の「New Token」をクリックしてアクセストークンを発行し、コピーしておきます。

これで事前準備は完了です。

環境設定

ここからは、Google Colaboratory上で作業を行います。

はじめに、画面上部のメニューから、「ランタイム」、「ランタイムのタイプを変更」と進み、「ハードウェアアクセラレータ」を「GPU」に変更しておきます。これにより、推論時にGPUを利用することができます。

続いて、必要となるライブラリのインストール、インポートを行います。 diffusersライブラリのバージョン0.3.0時点ではimage inpaintingに試験的にしか対応していないため、githubから直接コードを取得します。

# ライブラリのインストール
!pip install git+https://github.com/huggingface/diffusers.git
!pip install transformers scipy ftfy

# インポート
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw
from google.colab import files
import numpy as np
import torch
from torch import autocast
from diffusers import StableDiffusionInpaintPipeline

モデル設定

次に、コピーしておいたアクセストークンを利用して、Hugging Faceから学習済みモデルをダウンロードします。

# 学習済みモデルのダウンロード
access_token = 'ここにアクセストークンをペースト'
pipe = StableDiffusionInpaintPipeline.from_pretrained(
    'CompVis/stable-diffusion-v1-4',
    revision='fp16', 
    torch_dtype=torch.float16,
    use_auth_token=access_token
).to('cuda')

画像修復

それでは、画像修復を行っていきます。

image inpaintingでは、元の画像をベースとして、修復範囲を指定し、そこにどのような変更を加えたいかをprompt（プロンプト）と呼ばれる文字列で指示します。

今回は以下の写真をベースにします。

Reuters | Breaking International News & Views

はじめに、服をアロハシャツに変更する指示を与えてみます。顔から下の修復範囲が白、それ以外が黒となるようなマスク画像を作成し、プロンプトと合わせてモデルに渡します。

# 画像のアップロード
files.upload()

# ファイル名の設定
image_file = 'アップロードした画像のファイル名' 

# ベース画像の読込
original_image = Image.open(image_file).convert('RGB')
init_image = original_image.resize((512, 512))

# 修復領域の指定
x1 = 60  # 領域の左上のx座標
y1 = 170  # 領域の左上のy座標    
x2 = 512  # 領域の右下のx座標  
y2 = 512  # 領域の右下のy座標
mask = init_image.copy()
draw = ImageDraw.Draw(mask)
draw.rectangle([(x1, y1), (x2, y2)], outline=(0, 255, 0), width=3)
mask

# マスク画像の作成
draw.rectangle([(0, 0), (512, 512)], outline=(0, 0, 0), fill=(0, 0, 0))  # 全面を黒に
draw.rectangle([(x1, y1), (x2, y2)], outline=(255, 255, 255), fill=(255, 255, 255))  # 修復領域を白に

# プロンプトの設定
prompt = 'hawaiian shirt' 

# 画像の修復
 with autocast('cuda'):
        image = pipe(prompt=prompt, init_image=init_image, mask_image=mask)['sample'][0]
    image.save('test.png')