Google Colab で、お手軽 Text-to-Image

最近話題のお絵描きAIをGoogle Colabで試してみようという記事です。
ここで使うAIは、stability.aiの stable diffusion というモデルで、つい最近 CreativeML Open RAIL-Mライセンスの下、オープンソースでリリースされました。
https://stability.ai/blog/stable-diffusion-public-release

StableDiffusionに関する日本語記事は、深津 貴之さんの記事が面白いので是非こちらをお読みください。
https://note.com/fladdict/n/n13c1413c40de

stability.aiの記事で、Google Colabを用いて使う方法についての説明があり、試してみたところ本当に簡単に実行できてしまったので、その手軽さを紹介できればと思います。
実際に試したい場合は、元記事のGoogle Colabのノートブックが丁寧に説明しているので、そちらをご覧いただければと思います。

試してみる。

必要なもの

• Google アカウント
• Hugging Face アカウント

Hugging Faceとは

AIに関するプラットフォームです。
今回は、モデルやデータセット、デモアプリをオープンソースで公開出来るHugging Face Hubという場所から、stable diffusionのモデルを利用するために使います。

準備

GPUの使用

Google Colabからノートブックを開いて、そのノートブックがGPUを使用しているかを確認しましょう。

!nvidia-smi

次のように表示されたら、GPUの使用の設定が出来ていないので、メニューのランタイムからランタイムのタイプを変更をクリック、ハードウェア アクセラレータにGPUを選択して保存します。

NVIDIA-SMI has failed because it couldn't communicate with the NVIDIA driver. Make sure that the latest NVIDIA driver is installed and running.

その後、もう一度先ほどのコマンドを実行すると、使われているGPUのスペックが表示されます。

ライブラリインストール

!pip install diffusers==0.2.4
!pip install transformers scipy ftfy
!pip install "ipywidgets>=7,<8"

Google Colabの外部ウィジェット有効化

from google.colab import output
output.enable_custom_widget_manager()

Google CoalbからHugging Faceへログイン

アクセストークンをHugging Face側で作っておく必要があります。
アクセストークンのページ

from huggingface_hub import notebook_login

notebook_login()

パイプラインの作成

パイプラインの作成する前に、使用するモデルのリポジトリである https://huggingface.co/CompVis/stable-diffusion-v1-4 にてアクセス登録が必要でした。モデルを利用するための規約に同意するかの確認をするものです。
規約を読んでアクセス登録したらあとは遊ぶだけです。

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16, use_auth_token=True)

CUDAを利用するようにする。

pipe = pipe.to("cuda")

画像生成！

ここまで出来たら、あとは画像を作るだけです。
Google Colabが許す限り画像を生成できます。また、思っていた以上に画像の生成速度が速くて驚きました(15秒 ~ 30秒)

from torch import autocast

prompt = "a photograph of an astronaut riding a horse, Leica SUMMILUX 50mm f/1.4"
with autocast("cuda"):
  image = pipe(prompt)["sample"][0]

image

カメラのことは全く分かりませんが、落合陽一さんが良いと言っていたので、特に何も考えずにプロンプトにライカと付けています。これがあるだけでクオリティがかなり違います。
https://note.com/ochyai/n/n0f166ce4d713

とっても簡単に作れてしまったかと思います。
皆さんも是非遊んでみてください。