Stable Diffusionによる画像生成をGoogle Colaboratoryで実行する

この記事について

最近オープンソースで公開され話題のStable DiffusionをGoogle Colaboratory上で実行するまでの流れをまとめました。Stable Diffusionは任意のテキスト入力に対して写真のようなリアルな画像を出力するtext-to-imageのモデルです。これはDiffusion Modelという非平衡熱力学から発想を得た生成モデルの一種に基づいており、LAION-5Bという大規模なデータベースを活用して訓練を行ったものです。今回は、主に自然言語処理を対象にした大規模なオープンソースコミュニティーであるHugging Faceより訓練済みモデルをダウンロードします。

Hugging Faceへの登録

まずhttps://huggingface.co/よりHugging Faceにサインインする必要があります。

登録を終えたら「Settings > Access Token」より自分のアクセストークンを発行し、記録しておきます。

また、https://huggingface.co/CompVis/stable-diffusion-v1-4内の「Access Repository」をクリックすることで権限を得ましょう。またここで、モデルの作者にメールアドレスとユーザー名が共有されることに注意してください。

なお、モデルの使用にあたってはCreativeML OpenRAIL Licenseに準拠しなければならず、意図的に違法または有害な出力やコンテンツを作成・共有することの禁止や再配布のルールなどについての記載に同意する必要があります。

Google Colaboratoryでの実行

GPUの設定

Google Colaboratoryで新しいノートブックを作成した後、「ランタイム > ランタイムのタイプを変更」からハードウェアアクセラレータとして「GPU」を選択してください。

パッケージのダウンロード

以下のコマンドで、Stable Diffusionを始め必要なパッケージをダウンロードします。

!pip install --upgrade diffusers transformers scipy

重みのダウンロード

以下のの箇所を、Hugging Faceで取得したアクセストークンで置き換え、モデルの重みをダウンロードします。

download_weights.py

from diffusers import StableDiffusionPipeline

YOUR_TOKEN = "<User Access Tokens>"
model_id = "CompVis/stable-diffusion-v1-4"
pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=YOUR_TOKEN)
pipe.to("cuda")

画像生成

以下のを画像生成の入力となる文章やフレーズに置き換えて実行します。Google Colaboratory上に入力と同名のpng画像が保存されます。なお、入力テキストは英語にのみ対応しているようですが、日本語でもそれらしい出力が見られました。

generate.py

text = "<your original sentence>"
image = pipe(text)["sample"][0]
image.save("{}.png".format(text))

生成例

同じ文章（「雪国」の書き出し）から3つの画像を生成します。

generate_example.py

text = "The train came out of the long tunnel into the snow country"
num = 3
for i in range(num):
  image = pipe(text)["sample"][0]
  image.save("{}_{}.png".format(text,i))

素晴らしいクオリティとバリエーションに思えます。頭の中のイメージの具現化を容易にする技術の普及により、人々の創作活動がどう変化するのか興味深いです。