Stable Diffusionによる画像生成（image-to-image）

この記事では、Stable Diffusionを用いて画像とテキストから新たな画像を生成する方法を紹介します。

テキストから画像を生成するtext-to-image、画像の指定領域を修復するimage inpaintingの方法は、それぞれ以下の記事をご覧ください。

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を利用することができます。

続いて、必要となるライブラリのインストール、インポートを行います。

# ライブラリのインストール
!pip install diffusers==0.3.0 transformers scipy ftfy

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

モデル設定

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

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

画像生成

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

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

今回は以下のイラストを実写化する指示を与えてみたいと思います。

# 前処理関数の定義
def preprocess(image):
    w, h = image.size
    w, h = map(lambda x: x - x % 32, (w, h))
    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. * image - 1.

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

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

# プロンプトの設定
prompt = 'Photographic portrait of a detective, Sony α7, 35mm Lens, f1.8, golden hour, by Daniel F Gerhartz' 

# 生成画像数の設定
num = 10

# 初期画像の読込
original_image = Image.open(image_file).convert('RGB')
resized_image = original_image.resize((512, 512))
init_image = preprocess(resized_image)

# 画像の生成
for i in range(num):
    with autocast('cuda'):
        image = pipe(prompt=prompt, init_image=init_image)['images'][0]
    image.save(f'{i:03}.png')