StableDiffusionのcheckpointに任意のVAEをマージする

概要

Stable Diffusionのモデルは大きくVAE（Auto-Encoder）、Text Encoder、U-Netからなりますが、VAEは画像と潜在空間の変換を行っています。VAEを入れ替えることで出力される画像が変化し、たとえば特定のジャンル（アニメ絵など）がより高精細に出力される、といったことが期待できます。

そこでcheckpointのVAEだけを入れ替えて新しいcheckpointを作るスクリプトを作成しました。

生成環境によっては生成時に任意にVAEを組み合わせることができますが、あらかじめ組み合わせておくことでどのような環境でも使用可能になります。

スクリプト

以下をダウンロードして任意の場所に置いてください。StableDiffusionが動作する環境（PyTorchが動く環境）なら動作すると思います。

merge_vae.py

4.39 KB

ファイルダウンロードについて

ダウンロード

使用方法

以下のように入力してください（実際には1行で記述します）。

python merge_vae.py <VAEを入れ替えたいcheckpointファイル> 
    <入れ替えるVAEのcheckpointファイル> <入れ替えた後の出力先checkpointファイル>

入れ替えるVAEのcheckpointは、VAE単体（Waifu Diffusion 1.4のものなど）、通常のcheckpoint（VAEだけ抜き出します）、Diffusers形式のVAE単体（拡張子.bin、StableDiffusion 1.5のものなど）、いずれでも構いません。

サンプル

Waifu Diffusion v1.3のVAEを入れ替えて生成した例です（seed、prompt等の生成条件は同じ）。

Waifu Diffusion v1.3のVAE（SD 1.4と同じ）

Waifu Diffusion v1.4のVAEに入れ替え
色味が若干明るくなりまつ毛のあたりがすっきりしている

Stable Diffusion v1.5のVAEに入れ替え
やはりまつ毛のあたりが若干変化