Google ColabでJapanese Stable Diffusionを実行し、日本語の文章から画像を生成してみる #colaboratory

1. 背景

日本語の文章から画像を生成する、いわゆるジェネレーティブAIのJapanese Stable Diffusionが公開されたので、Colabで実行してみました。

2. アクセストークンの取得

以下サイトの「アクセストークンの取得」項から「モデルのライセンス条項等に同意する」項までを行い、トークンを取得します。トークンはコピーしておきましょう。

3. Colabにコードを記述

Colabに以下のコードを記載していきます。

※注意：Colabの環境は、GPUに設定してください
（ランタイム→ランタイムのタイプを変更から設定できます）

# Google Drive との連携
from google.colab import drive
drive.mount('/content/gdrive')

# japanese-stable-diffusionのinstall
! pip install git+https://github.com/rinnakk/japanese-stable-diffusion

ついで、以下のコードを記述します。先ほど取得したトークンを入力します。

import torch
from torch import autocast
from diffusers import LMSDiscreteScheduler
from japanese_stable_diffusion import JapaneseStableDiffusionPipeline
 
MODEL_ID = "rinna/japanese-stable-diffusion"
DEVICE = "cuda"
YOUR_TOKEN = #取得したトークンを入力#
 
# Use the K-LMS scheduler here instead
scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear",
                                 num_train_timesteps=1000)
 
pipe = JapaneseStableDiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.float32, scheduler=scheduler, use_auth_token=YOUR_TOKEN)
pipe = pipe.to(DEVICE)

ここまでで準備完了です。
最後に以下を入力します。image_numberの枚数だけ、指定したGoogleDriveのフォルダに生成した画像を格納します。なお、PROMPTは入力する文章になります。

PROMPT = '入力する文章'
image_number=4

for i in range(image_number):
 with autocast("cuda"):
    image = pipe(PROMPT, guidance_scale=7.5).images[0]
 image.save('/content/gdrive/MyDrive/image'+str(i)+'.png') #画像の保存先は適宜設定

参考

他のサイトなどを見ると、上記6行目が

    image = pipe(PROMPT, guidance_scale=7.5)["sample"][0]

となっていることがあります。上記で実行すると

KeyError: 'sample'

というエラーが出てしまいますので、

    image = pipe(PROMPT, guidance_scale=7.5).images[0]

とします。

4. 実行結果

今回はお題は PROMPT = "馬に乗る自由の女神の絵" として画像を作ってみました。

いずれもお題に合っているものを作ってくれています！
他にもお題に沿った画像を作ってくれました！

"山登りをする夫婦"

"スノーボードを楽しむ若者"

"公園で遊ぶ家族"

"ゲームの戦士のキャラクターデザイン"

５. まとめ

Japanese Stable Diffusionを使って、文章からさまざまな画像を作ってみました！日本語入力でここまでの画像を作れるとは・・・。rinna、恐るべしです。