【Stable Diffusion】Pythonでテキストから画像を生成しよう!

はじめに

この記事では、Stable Diffusionを用いて、テキストから画像を生成するプログラムをPythonで書きます。

また、日進月歩で新しい技術が登場することもあり、今回の記事で紹介した内容がすぐに古くなる可能性があります。そのため、今回紹介したコードが動作しなくなるなどの可能性も考えられますので、必ず公式のWebページ等もご参考いただけますと幸いです。

Stable Diffusionとは

Stable Diffusionとは、2022年8月に無償公開された画像描画AIです。
学習済みモデルが公開されており、Pythonから利用することが出来ます。
Stable Diffusionを用いることで、テキストでを与えることでそれに応じて画像が自動生成されます。
入力された文字列から画像を生成するタスクのことをText to Imageタスクと呼びます。
例えば、実際に「Cat flying in the sky」といった文字を入力することで以下のような画像を生成してくれます。

事前設定(Colabの設定)

今回は、Google Colaboratoryを用いて、コードを書いていきます。

まずランタイムをGPUに変更します。

GPUランライムの変更はこちらの記事をご参照ください。

【初学者必読】Google Colaboratory とは?使い方・便利な設定などをわかりやすく解説!

HuggingFaceのアカウント作成

Hugging Faceは、自然言語処理を対象にしたオープンソースの人工知能コミュニティーで、学習済みモデルやデータセットなどが公開されているサービスです。

HuggingFaceにアクセスし画面右上のSign Upよりアカウントを作成します。

メールアドレスとパスワードを入力すると以下のようなページに遷移します。

UsernameFull nameを入力します。

入力後、「I have read and agree with the Terms of Service and the Code of ConductCreate Account」にチェックを入れて、「Create Account」ボタンを押すことでアカウントの作成が完了します。

アカウント登録後のwelcomeページ

トークンの発行

https://huggingface.co/CompVis/stable-diffusion-v1-4
にアクセスし、ページ内の赤枠部分にチェックを入れて、Access repositryボタンを押して承認します。

承認後、HuggingFaceの画面右上のアイコンから[Settings]ページをクリックします。

左メニューに[Access Tokens]に移動し「NewToken」ボタンよりNameに任意の名前を入力し「Generate Token」よりトークン発行します。今回Nameには「Stable-diffusion-v1-4」と入力しています。

Stable Diffusionのインストール

PythonからStable Diffusionを利用するために、GPU設定したGoogle Colabのセルに以下のコマンドを入力し、実行ください。

!pip install diffusers==0.2.4 transformers scipy ftfy

DiffusersはHuggingFaceが開発しているライブラリです。画像や音声に対応した事前学習済みの拡散モデルが提供されており、モデルのダウンロードを簡単に行うことが出来ます。今回はDiffusersを用いて、「stable-diffusion-v1-4」を読み込みます。

モデルの読み込み

先ほど発行したトークンを変数access_tokenに指定し、以下のコードをColabに記載しコードを実行ください。

import torch
from torch import autocast
from diffusers import StableDiffusionPipeline

access_token = "発行したトークン"

# モデルのダウンロード
model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=access_token)
model.to("cuda")

これで事前準備は完了です。

テキストから画像を生成する

モデルの読み込みが無事出来ましたら、あとは文章をモデルに渡すことで画像が生成されます。

まずは、Colabのセルに生成した画像を保存しておくoutputsフォルダをmkdirコマンドで作成します。

!mkdir outputs

続いて文書から画像を生成するコードを書いていきます。

import matplotlib.pyplot as plt

num = 3 # 生成したい画像枚数
prompt = "Horse flying in the sky" # 生成させたいもととなる文章
 
for i in range(num):
    image = model(prompt)["sample"][0]
    # 保存
    image.save(f"outputs/{i}.png")

promptは生成させたいもととなる文章です。任意の文章に変えてみてください。
また、numという変数には生成する枚数を指定しています。
今回は3に設定しているため3枚の画像が生成されます。

画像は、outputsフォルダ内に保存されています。中身を確認すると画像が3枚生成されています。

うち1枚は以下のような画像でした。

またoutpusフォルダに生成された画像はmatplotlibを用いて以下のコードで描画可能です。

# outputsフォルダに保存された画像を描画
for i in range(num):
    plt.imshow(plt.imread(f"outputs/{i}.png"))
    plt.axis('off')
    plt.show()

AIの基礎から機械学習技術を効率よく身に付けたい方へ

AI Academy Bootcampの「オンデマンド動画+チャットサポートプラン」では、6か月35,000円にて、PythonプログラミングやAI基礎、統計学、機械学習、SQL、深層学習の基礎から、応用範囲まで幅広く学ぶことが出来ます。
独学でプログラミングを学習したが挫折してしまった方や、機械学習や深層学習に関してわからないことを講師にいつでも質問できる環境で学びたいという方は是非ご受講ご検討ください。