🐧

Stable Diffusion を Google Colabで実行する際に便利なコード

2022/08/27に公開

Python

雑記

Stable Diffusion

tech

初めに

Google Colaboratory で動かしていることを前提。

自分は Colab Pro にしているので、VRAMあたりで無料版の人は実行できないものがあるかも
その際は出力する画像数とかを減らせばいけると思います。

出力する先をGoogle Driveにしたり、画像をいっぱい生成したりできるようにまとめてる記事が
見当たらなかったので作成。

セットアップ

ほぼここを参照
https://note.com/npaka/n/ndd549d2ce556

本家の見た方が本質の理解には繋がるかも

# パッケージのインストール
!pip install diffusers==0.2.4 transformers scipy ftfy python-ulid

# トークン変数の準備
YOUR_TOKEN="huggingfaceから取得した自分のトークン"

from diffusers import StableDiffusionPipeline
from torch import autocast
import uuid

# StableDiffusionパイプラインの準備
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=YOUR_TOKEN)
pipe.to("cuda")
# NSFWの制限を外す
pipe.safety_checker = lambda images, **kwargs: (images, False)

画像生成

# 便利関数を準備

from PIL import Image
from torch import autocast
import uuid
from ulid import ULID

# 画像を繋げる関数
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

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

下のブロックを実行すると 9枚の画像と呪文のテキスト、3x3でまとめた画像が出力される。

num_images = 3

# ここに呪文を入れる
prompt_text = ""

prompt = [prompt_text] * num_images

# 一気に3枚以上出力するとVRAMの制限に引っかかるので、3回同じ処理をしてる。頭が悪い

with autocast("cuda"):
  images = pipe(prompt, num_inference_steps=150)["sample"]

with autocast("cuda"):
  images.extend(pipe(prompt, num_inference_steps=150)["sample"])

with autocast("cuda"):
  images.extend(pipe(prompt, num_inference_steps=150)["sample"])

ulid = ULID()

file_name = str(ulid)

# よくできた画像の呪文を忘れないようにファイル出力する
with open('/content/gdrive/MyDrive/googleColab/' + file_name + '.txt', 'w') as f:
  f.write(prompt_text)

# 3 x 3 の画像を出力する
grid = image_grid(images, rows=3, cols=3)
grid.save('/content/gdrive/MyDrive/googleColab/' + file_name + '.png')

# 作製された9枚の画像を一つずつ出力する
i = 1
for image in images:
  image.save('/content/gdrive/MyDrive/googleColab/' + file_name + '_' + str(i) + '.png')
  i = i + 1
grid

num_inference_stepsを変えると、若干クオリティに違いが出る。
大きければ大きいほど良いが、VRAM16GBだと、200程度が限界っぽい。（デフォルトは50）

成果物

私の編み出したpromptはこれ、num_inference_stepsは150くらいにしてます。

concept idea of the KAWAII girl with silver straight bangs and frilly dress for digital art on pivix fanbox, beautiful face, portfolio, thick coating painting, photorealistic, Fantastic, symmetrical, Arknights, Shadowverse, Azur Lane, WLOP, Makoto Shinkai, Yoh Yoshinari, Rembrandt Harmenszoon van Rijn style, Vtuber, 8k, 4K, trending on pixiv fanbox, trending on artstation, elaborate illustration, extremely high quality artwork, perfect art, brush stroke, clearly defined contours, insane detailed, dramatic lighting, very colorful, desaturated, beautiful face, tarot card, masterpiece, intente beautiful light, high contrast, blur

かなり成功確率が高くて、偶然を祈らなくてもいい感じの画像をバンバン出力してくれます。

追記：パラメータ弄って頑張ったらかなりいい感じの画像が出せました。
失敗作の方が多いですが、明らかに変なのは少ない感じです。

パラメータ

pipe(prompt, height=768, width=512, num_inference_steps=250)

promptは以下

concept idea of the KAWAII girl with platinum straight bangs and frilly dress for digital art on pivix fanbox, beautiful face, portfolio, thick coating painting, Bold line painting, photorealistic, Fantastic, symmetrical, Arknights, Shadowverse, Azur Lane, WLOP, Makoto Shinkai, Yoh Yoshinari, Rembrandt Harmenszoon van Rijn style, Vtuber, 8k, 4K, trending on pixiv fanbox, trending on artstation, elaborate illustration, extremely high quality artwork, perfect art, brush stroke, clearly defined contours, insane detailed, dramatic lighting, very colorful, beautiful face, full body, tarot card, magical, masterpiece, intente beautiful light, high contrast

こっちにここに書いてある以外のpromptと作例を列挙したものをまとめて載せています。
よかったら見てください。
https://note.com/6uclz1/n/nbc3d87d3e5b1

以上、参考になりましたら幸いです。

初めに

セットアップ

画像生成

成果物

Discussion