ptファイルのVAEをDiffusersで利用する方法

「VAEを利用して画質を上げたい・・・」
「DiffusersでptファイルのVAEをどうすれば読み込めるの？」

このような場合には、この記事の内容が参考になります。
この記事では、DiffusersでVAEを利用する方法を解説しています。

本記事の内容

Diffusersからpt形式のVAEを利用できない！！
変換スクリプトのダウンロード
pt形式からbin形式への変換
Diffusersからbin形式の読み込み

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

Diffusersからpt形式のVAEを利用できない！！

Diffusersを利用する場合、pt形式のVAEはそのままでは利用できません。
例えば、Anything v4.0の場合がそのケースになります。

andite/anything-v4.0 at main
https://huggingface.co/andite/anything-v4.0/tree/main

上記ページでは、以下の形式でVAEが公開されています。

ft-MSEのように次の形式であれば、Diffusersからも簡単に利用可能です。

この形式のVAEを利用する方法は、次の記事で解説しています。

したがって、Anything v4.0もbin形式（jsonファイル含む）で公開してくれればOK。
しかし、なかなかそこまで公開してくれるモデルはありません。

そのため、自分でpt形式からbin形式に変換する必要があります。
以下では、その方法を解説していきます。

変換スクリプトのダウンロード

GitHubページにあるDiffusers関連の変換スクリプト一覧
https://github.com/huggingface/diffusers/tree/main/scripts

本来であれば、上記ページにそのためのスクリプトがあるべきです。
しかし、見た感じではそのようなスクリプトは存在しません。
需要はあるはずなんですけどね・・・

でも、探せばあるものです。
心優しいエンジニアが、困っている人向けにスクリプトを提示しています。
~~https://huggingface.co/Linaqruf/anything-v3.0/discussions/66~~

追記 2023年1月26日
上記URLの内容は消えてしまったようです。

「convert_ldm_original_checkpoint_to_diffusers.py」を改良したようです。
と言っても、結構大幅に修正しているように見えます。

そのスクリプトは、上記掲示板でソースが公開されています。
それをコピペして、「convert_pt_to_bin.py」というファイル名で保存しましょう。
コピペが面倒な場合は、ここからダウンロードできます。

変換スクリプトが用意できたら、実際に変換してみましょう。
なお、PyTorchが動く環境ならスクリプトは動くと思います。

もし動かないなら、次の記事と同じ環境を用意してください。

pt形式からbin形式への変換

pt形式のVAEであれば、とりあえずは何でも構いません。
そもそも、そんなにVAE自体はそれほど多くは公開されていません。
むしろ、少ないです。

そんな状況において、取得できたVAEについては問題なく変換できました。
上記で紹介したAnything v4.0のVAEは、もちろん変換可能です。

変換する前に、スクリプトのヘルプを確認します。

> python convert_pt_to_bin.py -h
usage: merge_anything_vae.py [-h] --checkpoint_path CHECKPOINT_PATH [--original_config_file ORIGINAL_CONFIG_FILE] --dump_path DUMP_PATH

options:
  -h, --help            show this help message and exit
  --checkpoint_path CHECKPOINT_PATH
                        Path to the checkpoint to convert.
  --original_config_file ORIGINAL_CONFIG_FILE
                        The YAML config file corresponding to the original architecture.
  --dump_path DUMP_PATH
                        Path to the output model.

使い方は、簡単ですね。
実際には、次のコードで変換を行いました。

python convert_pt_to_bin.py  --checkpoint_path ..\..\ckpt\anything-v4.0.vae.pt --dump_path ..\..\model\anything4_vae

処理自体は、一瞬で終わります。
処理が完了すると、指定先にbin形式が生成されています。

config.jsonとbinファイルは、同じディレクトリに設置する仕様です。
最後に、このVAEを読み込んでみましょう。

Diffusersからbin形式の読み込み

VAEの読み込みは、すでに紹介した記事に従えば問題ありません。
再度、紹介しておきます。

唯一、悩むとすれば次の箇所です。

vae = AutoencoderKL.from_pretrained("./model/anything4_vae")

binファイルを指定するのではなく、ディレクトリのパスを指定します。
VAE読み込み部分をAnything v4.0のサンプルコードに組み入れると、次のようになります。

from diffusers import StableDiffusionPipeline
from diffusers.models import AutoencoderKL

model_id = "andite/anything-v4.0"

vae = AutoencoderKL.from_pretrained("./model/anything4_vae")
pipe = StableDiffusionPipeline.from_pretrained(model_id, vae=vae)
pipe = pipe.to("cuda")

prompt = "hatsune_miku"
image = pipe(prompt).images[0]

image.save("./hatsune_miku.png")

上記コードを実行すると、VAEを読み込んだ状態で画像が生成されます。
VAEの効果については、次の画像をご覧ください。