無料の画像生成AI「Stable Diffusion」をGoogle Colaboratory で動かす方法〈後編〉

2023年1月現在、トークンを使用した以下のプログラムは利用できなくなっています。新しい利用方法を別の記事で解説したので、そちらを参照してください。

前編では、Stable DiffusionをGoogle Colaboratoryで動かすためのコードを紹介しました。
後編では、複数枚の画像を同時につくる方法、画像をGoogleドライブに保存する方法、画像のサイズを変更する方法などについて紹介します。

プログラムコード
実行前に
1. フォルダの作成
2. 前編のプログラムの実行
Google Colabを使わなくても。

プログラムコード

最初に、紹介するプログラムをすべてまとめておきます。

前編のプログラムに追加して書いてください。

Googleドライブへの保存

画像の出力先をGoogleドライブの「生成画像」フォルダにします。
実行すると、Googleドライブと接続するか聞かれるので、許可してください。

# ドライブのマウント
from google.colab import drive
drive.mount('/content/drive')

# 出力先の指定
import os
os.chdir('/content/drive/MyDrive/生成画像')

複数画像の生成

1×3の形で3枚の画像を生成します。
数字を変えれば、2×2や3×3などにもできます。

from PIL import Image

def image_grid(imgs, rows, cols):
    assert len(imgs) == rows*cols

    w, h = imgs[0].size
    grid = Image.new('RGB', size=(cols*w, rows*h))
    grid_w, grid_h = grid.size
    
    for i, img in enumerate(imgs):
        grid.paste(img, box=(i%cols*w, i//cols*h))
    return grid

from torch import autocast
import datetime

#3枚の画像を1×3の形に並べて生成します。
num_rows = 1
num_cols = 3

prompt_single = "dancing cat"
prompt = [prompt_single] * num_cols

all_images = []
for i in range(num_rows):
  with autocast("cuda"):
    images = pipe(prompt)["sample"]
  all_images.extend(images)

# 画像を保存(ファイル名は「踊る猫_dancing cat_2022-2-22_22-22-22」などのようになる)
now = datetime.datetime.now(datetime.timezone(datetime.timedelta(hours=9)))
filename = "踊る猫"
i = 1
for image in all_images:
  image.save(f"{filename}_{prompt_single}_{now.strftime('%Y-%m-%d_%H-%M-%S')}_({i}).png")
  i += 1

grid = image_grid(all_images, rows=num_rows, cols=num_cols)
grid

※前編で「GPUが10GB以上の場合」のプログラムも紹介しましたが、それを使うとメモリの制限に引っかかります。ご注意を。

AIのパラメータ変更

学習のステップ数、画像のサイズなどを変更します。
image（複数画像の場合はimages）に代入するpipeの引数を、以下のように変えるとできます。

動画では分けて解説していますが、下のように一度にパラメータを設定することも可能です。

image = pipe(prompt,                     # キーワード
             height=512,                 # 画像の高さ
             width=512,                  # 画像の幅
             guidance_scale=7.5,         # 
             num_inference_steps=50,     # 画像生成に費やすステップ数
             generator=torch.Generator("cuda").manual_seed(0), # 乱数シードジェネレータ
             )["sample"][0]