Stable Diffusion 2をGoogle Colabでやってみた

澁谷直樹 @ キカベン

2022年12月8日 22:05

Stable Diffusion 2はTwitterでよく話題になっているAIモデルです。多くの人がうまく生成できた面白い画像をツイートしまくっています。

stability.aiによって開発された画像生成モデルになります。Google Brain東京のリサーチサイエンティストだったDavid Ha（Twitterのハンドルネーム@hardmaru)が最近移籍した先の会社でもあります。

今回は色々な画像を生成してみました。

Google ColabのGPUを使う準備

GPUを無料で使えるグーグルのColabのサービスを使います。

GoogleのColabを使うにはGoogleアカウントが必要です。また、ブラウザーはGoogle Chromeを勧めます。それでは、以下のリンクからGoogle Colabのアクセスしましょう。

「ノートブックを新規作成」をクリックして新しいノートブックを作ります。

適当なタイトルを付けておくと後で探すときに便利です。

まずGPUを使えるように設定します。編集メニューから「ノートブックの設定」を選びます。

そして、ハードウェアアクセラレータでGPUを選び保存をクリックします。

ノートブックを使う準備が整いました。

Huggingfaceのパイプラインを使う準備

ここからは、Huggingfaceにある説明に従います。

以下のコマンドをノートブックのセルで実行します。

!pip install --upgrade git+https://github.com/huggingface/diffusers.git transformers accelerate scipy

コピペしてプレイボタンをクリックするかSHIFT+ENTERで実行です。ちゃんとびっくりマーク（！）からコピーしてください。

インストールは40秒ほどかかります。

ちなみにプレイボタンを使うと自動でセルが追加されないので、「＋コード」をクリックしてセルを追加してください。

次に、PyTorchとStableDiffusionPipelineをインポートします。

import torch
from diffusers import StableDiffusionPipeline

特に出力もなくすぐに終わります。

次にモデルのIDを指定して、パイプラインを作成します。

from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler

model_id = "stabilityai/stable-diffusion-2"

# Use the Euler scheduler here instead
scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, revision="fp16", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

ここで訓練済みモデルのファイルがダウンロードされるのに1分少々かかります。

ここではサンプルに従ってEulerDiscreteSchedulerを使っています。スケジューラーは画像生成の際に加わえるノイズの発生手法の一つになります。ノイズが加わることで同じ文章でも毎回異なる画像が生成されます。

パイプラインはモデルを便利に扱えるように文章入力から画像生成までの一連の動作をまとめたものでHuggingfaceで提供されるモデルは通常パイプラインを用意してあります。

最後の行のpipe.to("cuda")でGPUを使うように指定しています。

以上でモデルの準備が整いました。

いろんなプロンプトで画像生成をやってみる

森で踊る犬と猫の漫画のイメージ

では、何らかのプロンプト（ヒントになる文章）を与えて画像を生成しましょう。

prompt = 'A cartoon image of a dog dancing with a cat in a forest.'
image = pipe(prompt).images[0]
image

「森で踊る犬と猫の漫画のイメージ」ですが、ちょっと想像してなかった奇妙な画像になりました。

画像生成には30-40秒ほどとかかります。

生成される画像はランダムなので同じプロンプトでも違うものになります。

なので、何度か繰り返すと気に入った面白いのが出てくるかもしれません。

ちなみに、日本語でのプロンプトでも一応動くようですが、かなり微妙な感じです。

prompt = '森で踊る犬と猫の漫画のイメージ。'
image = pipe(prompt).images[0]
image

以下は同じ日本語のプロンプトで２度目に生成したもの。

日本語よりはGoogle翻訳などで英語にしたものを使った方が良いかも（？）しれません。

馬に乗る宇宙服姿の写真

画像の大きさも指定が可能です。次のプロンプトはHuggingfaceのサンプルコードと同じものを使いました。

prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt, height=768, width=768).images[0]
image

サンプルにされるだけあって見事な出来栄えです。

宇宙映画「Stable Diffusion 2」と恐竜

うまくいくプロンプトを探すのも楽しいものです。この記事のタイトル画像もStable Diffusion 2で自動生成したものです。

prompt = 'A space movie title "Stable Diffusion 2" with dinosaur'
image = pipe(prompt).images[0]
image

サンタクロースが東京に来る

prompt = 'Santa is coming to Tokyo Tower on a starry snowy night'
image = pipe(prompt).images[0]
image

東京湾の海賊

prompt = 'Pirates in Tokyo Bay'
image = pipe(prompt).images[0]
image

シンギュラリティの後の世界

prompt = "The world after singularity"
image = pipe(prompt, height=768, width=768).images[0]
image

いかがでしょうか。やってるとキリがないですね。