本記事では、機械学習手法Stable Diffusionを用いて画像の特定の範囲をテキストに応じた画像に変更する画像修復を実行する方法をご紹介します。

Stable Diffusion

概要

Stable Diffusionは、拡散モデルによる画像合成モデルです。
自然言語で入力されたテキスト(prompt)から画像を生成するText to Imageタスクを実現します。

Latent Diffusionをベースとした本モデルは、非常に大規模なデータセットであるLAION-5Bを用いてトレーニングされています。

GoogleのImagenやOpenAIのGLIDE, DALLE-2など最新の画像生成モデルは社会に与える影響が大きいことなどを理由に学習済みモデルの公開を控えるケースが目立っていましたが、
Stable Diffusionでは、AIの平等な使用機会を与えることを目的に学習済みモデルの公開に踏み切っています。

なお、以下の記事では有償にはなりますが、Lattent DiffusionやGLIDEの技術解説やデモを紹介しています。よろしければご参照ください。
技術解説までは、無料で参照可能です。

本記事では上記手法を用いて、inpaitingタスクをGoogle Colaboratoryで実行していきます。

なお、テキストから画像を生成するtext2imgタスク、画像とテキストを入力し新たに画像を生成するimg2textは以下でご紹介しています。

デモ(Colaboratory)

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

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

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

環境セットアップ

初めにHuggingFaceのアカウントを作成します。
HuggingFaceにアクセスし画面右上のSignUpよりアカウントを作成します。
登録したメールアドレスに認証メールが届くのでメールに記載されたリンクにアクセスしてアカウント登録は完了です。

続いて、こちらのCompVis/stable-diffusion-v1-4にアクセスし記載の内容を確認の上承認します。

承認後、画面右上のアカウントのアイコンから[Settings]->[Access Tokens]に移動しNewTokenを発行し、メモしておきます。

以降の作業はGoogle Colaboratoryで行います。

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

まずGithubからコードをcloneします。

%cd /content

!git clone https://github.com/huggingface/diffusers.git
!sed -E -i "s/sample\(generator=generator\)/sample\(\)/" /content/diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py

次に、ライブラリをインストールします。

%cd /content

!pip install diffusers transformers scipy ftfy

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

import matplotlib.pyplot as plt
from PIL import Image, ImageDraw

import torch
from torch import autocast

%cd /content/diffusers

from src.diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint import StableDiffusionInpaintPipeline

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

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

続いて、先ほど発行したアクセストークンを使用して、モデルをダウンロードします。

%cd /content/diffusers
access_tokens="先ほどメモしたAccessToken" # @param {type:"string"}

# load pretrain model

%cd /content/diffusers

device = "cuda"
pipe = StableDiffusionInpaintingPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    revision="fp16", 
    torch_dtype=torch.float16,
    use_auth_token=access_tokens
).to(device)

マスク画像生成

それでは、promptと画像を設定し指定範囲を画像修復していきます。
画像修復は、マスク画像で指定された白色領域のみに対して画像生成することで実現します。

まず、入力画像の取得と、マスク画像の生成を行います。

img_w, img_h = 768, 512

%cd /content/diffusers

init_image = Image.open(filename)
if '.png' in filename:
  init_image = init_image.convert('RGB')
init_image = init_image.resize((img_w, img_h))

plt.imshow(init_image)
plt.axis('off')
plt.show()

以下の画像を使用します。

続いて、画像修復範囲を矩形で指定します。

x1 = 400  #@param {type:"integer"}
y1 = 0    #@param {type:"integer"}
x2 = 650  #@param {type:"integer"}
y2 = 300  #@param {type:"integer"}

%cd /content/diffusers

temp = Image.open(filename)
if '.png' in filename:
  temp = temp.convert('RGB')
mask_img = temp.resize((img_w, img_h))
draw = ImageDraw.Draw(mask_img)

# 矩形の設定
rectcolor = (255, 0, 0)
linewidth = 4
draw.rectangle(
    [(x1, y1), (x2, y2)], # 左上x,y 右下x,y
    outline=rectcolor, width=linewidth,
    )
mask_img

出力結果は以下の通りです。

指定された矩形に従ってマスク画像を生成します。

# mask画像生成

# 一度すべて黒く塗りつぶし
rectcolor = (0, 0, 0)
draw.rectangle(
    [(0, 0), (img_w, img_h)],
    outline=rectcolor, fill=(0, 0, 0)
    )
# 先ほど設定した矩形箇所白く塗りつぶし
rectcolor = (255, 255, 255)
draw.rectangle(
   [(x1, y1), (x2, y2)],
    outline=rectcolor, fill=(255, 255, 255)
    )
mask_img

以下のマスク画像が生成されます。

白色の領域が修復範囲です。

Inpainting

入力画像、マスク画像がそろったのであとはモデルに入力するのみです。

promptにganesha faceを設定して実行してみます。

prompt = "ganesha face" #@param {type:"string"}

generator = torch.Generator(device).manual_seed(12)

with autocast("cuda"):
  images = pipe(
      prompt=prompt,
      init_image=init_image,
      mask_image=mask_img,
      strength=0.75,
      guidance_scale=7.5,
      generator=generator,
  )["sample"]

images[0].save("inpainting_01.png")

出力結果は以下の通りです。

矩形範囲のみガネーシャの顔に修復することができました。

まとめ

本記事では、Stable Diffusionを用いてinpaintingを行う方法をご紹介しました。
かなり高精度なモデルのため悪用厳禁であることは言うまでもありません。

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

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

参考文献

1.  論文 - High-Resolution Image Synthesis with Latent Diffusion Models

2. GitHub - CompVis/stable-diffusion