Stable DiffusionのモデルをU-Netの深さに応じて比率を変えてマージする
概要
Stable Diffusionでは、複数のモデルの重みをマージすると中間的な出力が得られることが知られています。
以下の比較結果が大変参考になります。
またマージ用のスクリプトは以下で公開されています。
また(ざっくりとした私の理解では)Stable Diffusionは大きくText Encoder (CLIP)、Denoising Auto Encoder (U-Net)、Auto Encoderからなります。このうち、U-Netが画像をノイズから画像を生成する部分担当しています。
U-Netは広く使われているネットワーク構造で、構造がその名の通りU字型(V字型)をしています。
またU-Netに限らず一般的に、ニューラルネットワークでは浅い層が細かいディテール部分を認識し、深い層が全体の構造を認識するといわれています。
そこで、U-Netの深さに応じてマージ比率を変えてマージすると、単純にマージしたのとは異なる結果が出るのではないか、と思い試してみました。
具体的には、U-Netの入力側、出力側にそれぞれ12個(12階層)、中間に1個のブロックがありますので、それぞれの層でマージモデルの比率を変えます。
また今回は比較のため、Text Encoder、Auto Encoderは、どちらか片方を100%として使います。
比率を指定してマージするスクリプトは記事の末尾に添付します。
マージするモデル、プロンプト
今回はStable Diffusionモデルと、Waifu Diffusion 1.3をマージして試してみます(わかりやすいように傾向が異なる二つのモデルを選びました)。
プロンプトはとりあえず適当に以下を用いました。
masterpiece, best quality, beautiful anime girl, school uniform, strong rim light, intense shadows, highly detailed, cinematic lighting, taken by Canon EOS 5D Simga Art Lens 50mm f1.8 ISO 100 Shutter Speed 1000
seed : 1~4
sampler: k_euler_a
steps : 80
scale : 7.0
negative-prompt: lowres, bad anatomy, bad hands, error, missing fingers, cropped, worst quality, low quality, normal quality, jpeg artifacts, blurry
生成結果
以下が生成結果になります。
単体モデル
まずはStable Diffusionです。

Waifu Diffusionです。

単純マージモデル
50:50でマージしました。それらしい結果が出ています。

比率付きマージモデル
まず浅い層をStable Diffusion 100%、最も深い層をWaifu Diffusion 100%として、間を均等に割り振りマージしたモデルです。Text Encoder、Auto EncoderはSDのものです。

意外にアニメ顔になっています。アニメ顔にするかどうかは、U-Netの深い層が担っているようです。
次に浅い層をWaifu Diffusion 100%、深い層をStable Diffusion 100%として、間を均等に割り振りマージしたモデルです。Text Encoder、Auto EncoderはWDのものです。

どちらかというとリアル寄りの、50:50とも違う、興味深い結果になりました。塗りの傾向はWDっぽい感じもあります。
マージ用スクリプト
冒頭のeyriewow氏のスクリプトを改造したものです。PyTorchがインストールされた環境で実行してください(氏のページを参考にしてください)。
コマンドラインオプションは以下の通りです(--helpも参照してください)。指定例はその次に記述します。
model_0
マージ元モデル1
model_1
マージ元モデル2
--base_alpha
Text Encoder、Auto Encoderでモデル1を使う割合。1.0でモデル1、0.5で50:50でマージします。
--device
モデルをどこにロードするか。デフォルト"cpu"なので通常は特に指定する必要はありません。"cuda"を指定するとGPUへ読み込みます。
--weights
U-Netの各層に対応した比率を、カンマ区切りの数値、25個で指定します。順に入力層の12個のブロック、中間の1ブロック、出力層の12個のブロックの、モデル1側の比率になります。1.0でモデル1が100%、0.5で50%、0で0%になります。
省略するとbase_alphaに指定した比率での単純なマージになります。
--verbose
指定するとU-Netの各層の比率を表示します。
コマンドラインの記述例は、たとえば次のようになります(モデル1を浅い層で100%、モデル2を深い層で100%、間を均等に割り振る場合)。
python merge_block_weighted.py model1.ckpt model2.ckpt
--base_alpha 1.0
--weights 1,0.9166666667,0.8333333333,0.75,0.6666666667,0.5833333333,0.5,
0.4166666667,0.3333333333,0.25,0.1666666667,0.0833333333,0,
0.0833333333,0.1666666667,0.25,0.3333333333,0.4166666667,0.5,
0.5833333333,0.6666666667,0.75,0.8333333333,0.9166666667,1.0
実際には1行で入力してください。
まとめ
単純なマージとは異なる結果になりました。浅いほうがディテールに影響する、といった感じでもなく、なかなか興味深いです。
今回は単純に深さで比率を変えましたが、「入力層をモデル1に、出力層をモデル2にする」といったことや、それぞれの比率を0/100でなく30/70にする、Text EncoderやAuto Encoderのマージ比率など、いろいろと方法は考えられます。お試しいただければ幸いです。