Diffusersを使ってColabでAIに絵をかいてもらいたい

AIを使って何かをしたいと思い、タロットカード占いでも作ってみようかと思った。
しかし、私には絵心がない。
タロットカードの素材なんかもなかなか転がっていない。

・・・AIに書いてもらえばいいのでは。
というわけでやってみたメモ。

環境
Google Colaboratory

GoogleColaboratoryについて

以前まではWebUIが使えたそうなのですが、無料の範囲では使用できなさそうです。
そのため、DiffuserのPipelineを使用して生成していきます。
また、GPUの使用時間も8時間程度と制限があったり、しばらく放置しているといつの間にか切断されていていたりします。
注意して使っていきましょう。

Diffuserのインストール

まずはDiffuserのインストールをするためのコードを書いていく。

!pip install diffusers transformers scipy ftfy accelerate

これのみです。
GPUへの接続は時間がもったいないので、すべてのコードを書き込んでから行います。

LoRAを読み込めるようにする

こちらのGitHubを参考にした別の記事からの拾いものですが・・・。
コードに起こすと下記です。やってみておかしかった場合は上記のGitHubで更新があるかと思いますので、参照に修正してみてください。

import torch
from safetensors.torch import load_file

def load_safetensors_lora(pipeline, checkpoint_path, LORA_PREFIX_UNET="lora_unet", LORA_PREFIX_TEXT_ENCODER="lora_te", alpha=0.75):
    # load LoRA weight from .safetensors
    state_dict = load_file(checkpoint_path)

    visited = []

    # directly update weight in diffusers model
    for key in state_dict:
        # it is suggested to print out the key, it usually will be something like below
        # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight"

        # as we have set the alpha beforehand, so just skip
        if ".alpha" in key or key in visited:
            continue

        if "text" in key:
            layer_infos = key.split(".")[0].split(LORA_PREFIX_TEXT_ENCODER + "_")[-1].split("_")
            curr_layer = pipeline.text_encoder
        else:
            layer_infos = key.split(".")[0].split(LORA_PREFIX_UNET + "_")[-1].split("_")
            curr_layer = pipeline.unet

        # find the target layer
        temp_name = layer_infos.pop(0)
        while len(layer_infos) > -1:
            try:
                curr_layer = curr_layer.__getattr__(temp_name)
                if len(layer_infos) > 0:
                    temp_name = layer_infos.pop(0)
                elif len(layer_infos) == 0:
                    break
            except Exception:
                if len(temp_name) > 0:
                    temp_name += "_" + layer_infos.pop(0)
                else:
                    temp_name = layer_infos.pop(0)

        pair_keys = []
        if "lora_down" in key:
            pair_keys.append(key.replace("lora_down", "lora_up"))
            pair_keys.append(key)
        else:
            pair_keys.append(key)
            pair_keys.append(key.replace("lora_up", "lora_down"))

        # update weight
        if len(state_dict[pair_keys[0]].shape) == 4:
            weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32)
            weight_down = state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32)
            curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3)
        else:
            weight_up = state_dict[pair_keys[0]].to(torch.float32)
            weight_down = state_dict[pair_keys[1]].to(torch.float32)
            curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down)

        # update visited list
        for item in pair_keys:
            visited.append(item)

    return pipeline

Pipelineを作成する

描画するためのPipelineを作成します。
生成モデルはHuggingFaceのものを使用します。

https://huggingface.co/

登録しなくても使用できそうな感じはするのですが、ちゃんと登録しておいたほうがお気に入りの整理もできていいかもしれません。
使いたいモデルを見つけたら名前をコピーしましょう。
私は Lykon/NeverEnding-Dream が好きでよく使います。

import torch
from diffusers import StableDiffusionPipeline

model_id = "Lykon/NeverEnding-Dream"
device = "cuda"

pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    safety_checker=None,
)
#pipe = load_safetensors_lora(
#    pipe,
#    'lora_name.safetensors'
#)
pipe = pipe.to(device)

safety_checkerは性的な画像を黒く塗りつぶして表示しなくする機能です。
ただ、なぜかわからないけど黒く塗りつぶされることが結構あるので、よほど見たくない限りはNoneにしてオフしておいたほうが原因がわかっていいような気がします。
ただし、HuggingFaceからちょっと警告はされます。
LoRAを使用したい場合はコメントアウトを消してLoRAを読み込んでください。
こちらは生成モデルと違い、Driveまたは作業フォルダにアップロードして使うことになります。

生成プロンプト

最後に生成プロンプト入力と画像の生成を行います。

prompt = "(masterpiece), (extremely intricate:1.3),  8k, best quality, ((photorealistic)), ultra detailed, detailed background, ((raw photo)), nature"
n_prompt = "low quality, worst quality, nsfw"
# 画像の生成、保存、表示
image = (pipe(prompt, negative_prompt=n_prompt, guidance_scale=7.5, num_inference_steps=50, height=512, width=512).images[0])
image[cnt].save(f"man{cnt}.png")
image

promptに生成したい画像の特徴、n_promptに生成したい画像に取り込みたくない特徴を入力します。
width, heightで画像サイズを指定していますが、基本的にはこの512x512がAIの生成に一番いいといわれています。生成してから高画質化するなり、周りを書き足してもらうなりで大きな画像にしていくようです。
とりあえず、すごく抽象的なものを作成してみます。それぞれに単語の意味は試しながらお勉強してください。呪文で調べるといろいろと出てきます。

GPUランタイムを使用する

それでは準備が整ったのでGPUを使用していきます。
メニューバーからランタイムのタイプを変更を選択します。

T4 GPUを選択して保存します。
これでGPU使用の準備が整いました。

画像を生成する

右上のランタイムの状態が下の画像のようになっていればOKです。

上から順番にコードを実行していきます。
パイプラインの生成に一番時間がかかるような印象です。
最後まで実行すると画像が1つ表示されるはずです。

心が洗われますね。
画像は右クリックで保存ができます。複数枚生成するときはfor文で回して保存してから表示するといいかと思います。

とてもお手軽でした。使用時間の限度はあるものの、ちょっと試してみるにはいいですね。

タロットカードも生成してみた

ちなみに、LoRAを入れてタロットカードも生成してみました。
こちらは生成モデルに gsdf/Counterfeit-V3.0 とLoRAに Tarot Cards (Rider-Waite) を使用しています。（正確にはちょっと違うんですが・・・。）
何のカードだと思いますか？