Stable DiffusionのモデルをU-Netの深さに応じて比率を変えてマージする

Kohya S.

概要

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です。

Stable Diffusion 100

Waifu Diffusionです。

Waifu Diffusion 100

単純マージモデル

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

Stable Diffusion 50: Waifu Diffusion 50

比率付きマージモデル

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

SD 100->0->100、WD 0->100->0

意外にアニメ顔になっています。アニメ顔にするかどうかは、U-Netの深い層が担っているようです。

次に浅い層をWaifu Diffusion 100%、深い層をStable Diffusion 100%として、間を均等に割り振りマージしたモデルです。Text Encoder、Auto EncoderはWDのものです。

SD 0->100->0、WD 100->0->100

どちらかというとリアル寄りの、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のマージ比率など、いろいろと方法は考えられます。お試しいただければ幸いです。

この記事が気に入ったらサポートをしてみませんか？