[Stable Diffusion] AIでテキストと画像から新たな画像を生成する [img2img]

2022年8月27日土曜日

Artificial Intelligence

本記事では、機械学習手法Stable Diffusionを用いてテキストと画像から新たな画像を生成するimg2imgを実行する方法をご紹介します。

アイキャッチ

Stable Diffusion

概要

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

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

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

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

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

なお、テキストから画像を生成するtext2imgタスク、特定の範囲の画像を修復するInpaintingタスクは以下でご紹介しています。

デモ(Colaboratory)

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

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

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

環境セットアップ

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

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

承認

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

Settings

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

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

「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」を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_img2img.py

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

%cd /content

!pip install diffusers transformers scipy ftfy

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

import matplotlib.pyplot as plt
from PIL import Image

import torch
from torch import autocast
import numpy as np

%cd /content/diffusers

from diffusers import StableDiffusionImg2ImgPipeline

最後に前処理関数を定義しておきます。

def preprocess(image):
    w, h = image.size
    w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
    image = image.resize((w, h), resample=Image.LANCZOS)
    image = np.array(image).astype(np.float32) / 255.0
    image = image[None].transpose(0, 3, 1, 2)
    image = torch.from_numpy(image)
    return 2.0 * image - 1.0

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

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

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

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

# load pretrain model

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

Image 2 Image

それでは、promptと画像を設定し新たな画像を生成していきます。

%cd /content/diffusers
filename = 'test_01.jpg'

!rm -rf {filename}

!wget -c https://www.pakutaso.com/shared/img/thumb/SAYA160312500I9A3721_TP_V4.jpg \
      -O {filename}
      
init_image = Image.open(filename)
if '.png' in filename:
  init_image = init_image.convert('RGB')
resize_image = init_image.resize((768, 512))
input_img = preprocess(resize_image)

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

以下の画像をモデルに入力します。

input image
prompt = "detailed CG art, woman at cafe" #@param {type:"string"}

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

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

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

plt.imshow(plt.imread("img2img_01.png"))

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

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

result

入力画像を構図をそのままにCGのような画像が生成されました。
使い方次第で様々な用途が考えられそうです。

まとめ

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

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


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

参考文献

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

2. GitHub - CompVis/stable-diffusion

AIで副業ならココから!

まずは無料会員登録

プロフィール

メーカーで研究開発を行う現役エンジニア
組み込み機器開発や機会学習モデル開発に従事しています

本ブログでは最新AI技術を中心にソースコード付きでご紹介します


Twitter

カテゴリ

このブログを検索

ブログ アーカイブ

TeDokology