ControlNetの軽量版？つくった

2024年3月7日 21:29

　ControlNetはすごいですが、UNetの3分の1～2分の1くらいの大きさのモデルが余計に必要です。今回は軽量なモデル＋LoRAだけでControlNetみたいなことをやってみます。

ComfyUI用実装

https://github.com/laksjdjf/LoRTnoC-ComfyUI

あーきてくちゃあ

　ControlNetは基本的にUNetのIN+MIDブロックのコピーですが、cannyなどの参照画像を潜在変数と同じ形状（4チャンネルでサイズが8分の1）にするinput hint block及びUNetとControlNetをつなげるzero convがあります。zero convはトレーニングの初期段階では重みがゼロになっており、UNetへの情報を遮断することで初期の学習を安定させる役割があります。
　今回私が試してみた設計は、input hint block及びLoRAを組み合わせたものになります。input hint blockの出力をUNetのconv_in(最初のモジュール)のチャンネルと同じ320にして、足し合わせます。あとはLoRAに任せるという感じです。LoRAは初期化時は出力が0になるよう重みを調整されているので、ちょうどzero convと似たような働きになります。

　一応この設計の根拠としてSD-2-depthがあります。これはdepth画像をリサイズして潜在変数とチャンネル方向に結合した5チャンネルをUNetへの入力とするモデルみたいです。おそらくconv_inの入力を1チャンネル分拡張して、あとはSD2のコピーから学習をしてると思います。そう考えるとconv_inの拡張＋LoRAだけでいけそうですよね。depthと違ってcannyなどはリサイズに弱そうなのでControlNetで使われているinput hint blockを使ってサイズを8分の1にすることにしました。UNetのconv_inは変えずにinput hint blockの出力チャンネルをconv_inの出力と合わせることで、出力同士の加算でいいようにしています。ComfyUI的にもこれがやりやすい。DeepShrinkのために実装された機能が使えます。

名前

LoRTnoCという略語の方を先に考えたので無理やりあてはめようとしています。
LoRA with hint block of ControlNet
がいい感じかな。

結果

　githubにいっぱいあげた。

長所

生成時間がほとんど増えない

　生成ループの最初に軽量のinput hint blockを計算し、LoRAをモデルへマージすれば同じです。他の手法と違って計算量がステップ数に依存しないので、1ステップでも1000ステップでも計算時間の増加は同じ大きさになります。その辺はT2i-Adapterと似てますね。

ファイルサイズが小さい

　SDXL用がfp16で8MBちょっとになります。

実装がらくらくちんちん

　LoRA使うのでその部分の実装は使いまわせます。あとはconv_inの出力にinput hint blockの結果を加算するだけであり、ComfyUI上では簡単（ということはForgeでも簡単なんじゃね？）

短所

表現力が低い

　他の手法と違って参照画像の情報を取り込むのが入力部分(conv_inの出力）1個だけになります。別に他の部分にも入力したっていいんですが、それはもうT2i-Adapterやんと思ってやめました。
　poseが特に微妙ですね。。。

柔軟性が低い

　LoRAを使うので特定範囲の適用や時刻の制限がやりづらいです。また複数適用の結果も微妙なことになりそうです。

LoRAとばっちんぐしそう

　LoRAを使うので他のLoRAとの相性問題があります。

スタイル等も学習しちゃいそう

　LoRAを使うのでスタイル等学習データの学習して欲しくない部分まで学習しちゃいます。

結論

　短所をあげるときりがないくらい微妙なところも多いですが、正直ControlNetなんて使わないので応用できるかどうかなんてどうでもいいや。