【Stable Diffusion】ckptをDiffusersで読み込む

「Stable Diffusionでckptファイルをモデルとして利用したい」
「ckptをDiffusersで読み込めるモデルに変換したい」

このような場合には、この記事の内容が参考になります。
この記事では、ckptをDiffusersでモデルとして読み込む方法を解説しています。

本記事の内容

Diffusers用モデルへのckptの変換
変換に必要なモノ
変換方法
変換したモデルの読み込み
「No module named ‘pytorch_lightning’」が出る場合の対策

それでは、上記に沿って解説していきます。

Diffusers用モデルへのckptの変換

画像生成AIのStable Diffusionは、公開以降の勢いがスゴイです。
その勢いに伴って、ファインチューニングしたモデルが次々に公開されています。

例えば、Waifu Diffusionです。
Waifu Diffusionについては、次の記事で説明しています。

他には、ポケモンのキャラをファインチューニングしたSD Pokémon Diffusersがあります。

これらは、Diffusersがインストール済みであれば簡単に利用できます。
単純な場合には、モデル指定の値を変更するだけです。

そんな中、ckptファイルがモデルとして公開されていることがあります。
ただ、ckptファイルはDiffusersで簡単に読み込むことができません。

例えば、HentaiDiffusionではckptがモデルとして公開されています。

Diffusersで読み込むためには、ckptを変換する必要があります。
変換すれば、Diffusersでモデルとして読み込めます。

以上、Diffusers用モデルへのckptの変換を説明しました。
次は、変換に必要なモノを説明します。

変換に必要なモノ

ckptをDiffusers用モデルに変換するには、以下が必要です。

Diffusers
Transformers
OmegaConf
Diffusersソース一式（GitHub）
Stable Diffusionの設定ファイル（v1-inference.yaml）

それぞれを以下で説明します。

DiffusersとTransformers

Diffusersライブラリのインストールが必要です。
Diffusersについては、次の記事で解説しています。

上記記事では、Transformersのインストールについても説明しています。
Transformers自体の詳細は、次の記事で解説しています。

OmegaConf

設定ファイルの読み込みに必要となります。
ここで言う設定ファイルとは、YAML形式のファイルです。

次のコマンドでインストールできます。

pip install OmegaConf

Diffusersソース一式（GitHub）

変換用のスクリプトが必要になります。
そのスクリプトが、GitHub上のDiffusersプロジェクトに含まれています。

ソース一式は、以下のコマンドで取得します。

git clone https://github.com/huggingface/diffusers

Windowsの場合は、Git for Windowsをインストールしておきましょう。
Windowsでもgitコマンドが利用できます。

Stable Diffusionの設定ファイル（v1-inference.yaml）

v1-inference.yamlは、Stable Diffusionの設定ファイルになります。
以下のページからダウンロードできます。

stable-diffusion/v1-inference.yaml at main · CompVis/stable-diffusion · GitHub
https://github.com/CompVis/stable-diffusion/blob/main/configs/stable-diffusion/v1-inference.yaml

まとめ

最低限で必要なモノを説明してきました。
早い話、Stable Diffusionが動く環境が必要です。

それにプラスアルファで必要なモノがあるということですね。
この場合、すでにStable Diffusionが動いている環境をそのまま利用するかどうかとなります。

個人的には、分けておいた方がよいと思います。
変換用の仕組みは、別のPython仮想環境に設けることを推奨します。

以上、変換に必要なモノについて説明しました。
次は、変換方法を説明します。

変換方法

ckptをDiffusers用モデルに変換していきます。
まずは、ckptの取得からです。

今回は、Stable Diffusionの標準モデルであるckptをダウンロードします。
ckptは、以下のページからダウンロードできます。
https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/tree/main

ここでは、「sd-v1-4.ckpt」を選択。
4GB以上あるので、ダウンロードに時間はかかります。

保存場所は、各自の都合の良い場所でOKです。

「diffusers」は、GitHubから取得したDiffusersソース一式です。
その付近に関係のあるディレクトリを作成しました。
今回は、「ckpt」と「model」を新規作成しています。

ダウンロードしたckptは、「ckpt」ディレクトリに保存しています。
あと、「v1-inference.yaml」（設定ファイル）は次の場所に設置しましょう。

「convert_original_stable_diffusion_to_diffusers.py」が、利用する変換用のスクリプトです。
このスクリプトと同じ階層に「v1-inference.yaml」を設置します。

ここまでの作業で準備は整いました。
では、「convert_original_stable_diffusion_to_diffusers.py」がある場所まで移動します。

ここで実行するのは、以下のコマンドです。
ckptのパスとモデルの出力を指定しています。

python .\convert_original_stable_diffusion_to_diffusers.py --checkpoint_path "..\..\ckpt\sd-v1-4.ckpt" --dump_path "..\..\model\sd"

上記を実行して、処理が進んだらラッキー！！
次のようなエラーが出たら、後述の対策をご覧ください。

ModuleNotFoundError: No module named 'pytorch_lightning'

処理が進んでラッキーな場合でも、しばらく待たされます。
指定したモデルの出力先に、ディレクトリが作成されている様子を確認できます。

..\..\model\sd
..\..\model\sd\vae
..\..\model\sd\vae
..\..\model\sd\text_encoder
..\..\model\sd\text_encoder
..\..\model\sd\tokenizer
..\..\model\sd\unet
..\..\model\sd\unet
..\..\model\sd\scheduler
..\..\model\sd\safety_checker
..\..\model\sd\safety_checker
..\..\model\sd\feature_extractor

処理が完了したら、確認しましょう。
最終的には、次のような構成になっています。

ここまでの作業により、ckptをDiffusers用モデルに変換できました。

以上、変換方法を説明しました。
次は、変換したモデルの読み込みを説明します。

変換したモデルの読み込み

基本的には、モデルに合わせた読み込みが必要となります。
今回の「sd-v1-4.ckpt」であれば、次のコードでモデル読み込みが可能です。

import torch
from diffusers import StableDiffusionPipeline
from torch import autocast

MODEL_PATH = "./model/sd"
DEVICE = "cuda"

pipe = StableDiffusionPipeline.from_pretrained(MODEL_PATH, torch_dtype=torch.float16)
pipe.to(DEVICE)

prompt = "a dog painted by Katsuhika Hokusai"

with autocast(DEVICE):
    image = pipe(prompt, guidance_scale=7.5)["sample"][0]
    image.save("test.png")

from_pretrained関数は、モデル名かモデルパスを設定します。
モデル名であれば、みなさんも普通に指定していることでしょう。