📖

[Stable Diffusion] Depth-to-Imageモデルを学習なしで特定のドメインに適応させる

2023/01/19に公開

ディープラーニング

Stable Diffusion

tech

はじめに

こんにちは。

早速本題に入ります。

[2023/01/20追記]
この記事で説明しているDepth-to-Imageモデルのドメイン適応と、適応したモデルのAUTOMATIC1111さんのWebUIでの使用が行えるColab Notebookを公開しました。

こんな感じで使えます

Depth-to-Imageモデルの説明

StabilityAIが公開したStable Diffusion V2系のモデルの中に、depthモデルというものがあります。

このモデルができることは、一般的なStable Diffusionのimg2imgができることと基本的には同じで、画像とテキストを入力とし、入力された画像スタイルをテキストの指示に従って変換する、というものです。

まずは、一般的なモデル（ここではWaifu Diffusion V1.4 Epoch1を使用しています）のimg2img結果を見てみましょう。

変換前

入力画像
入力テキスト: masterpiece, 1boy, anime

変換後

.
.
.

上の例を見てもらうとわかるかと思いますが、従来的なimg2imgは課題を抱えています。
それは、入力画像のスタイル変換（上の例だと実写→アニメ調への変換）をしっかりと行うためには、入力画像をある程度破壊的に変更して変換してもらう必要があり、その場合、例のように影がヨロイか何かに変化するような、入力画像から諸々大きくかけ離れた出力しか得られないことがあります。

だからといって入力に忠実な変換をしようとすると、うまくスタイル変換ができません。

入力に忠実な変換をしようとした場合(Denoising Strength=0.35)

そういう感じではない

アニメ調に変換したかったはずが、なんというかリアル系のイラストに変換されてしまいました。入力画像にあるオブジェクトの位置関係や形状について破壊的変更をしないよう画像を変換する場合、現実の写真から大きく離れたスタイルであるアニメ調には辿り着けないということです。

以上に書いたように、現状のimg2imgが抱える取り回しの悪さに対処すべくStabilityAIが公開したのが、画像の深度情報を保ちつつスタイル変換(depth-to-image, depth2img)を行うことができる、depthモデルです。

このモデルによるdepth2imgは、一般的なStable Diffusionと異なり、UNetにノイズを載せた画像（の特徴量）だけを入力するのではなく、別の深層学習モデルによって推測された画像の深度情報も入力し、UNetに対して直に画像内のオブジェクトの場所や遠近を教え込みます。UNetは深度情報と一貫性を保った出力を作成するように学習されているため、画像に対して強めのスタイル変換をかける場合でも、画像内のオブジェクトがどっかに行ったり変なものが追加されたりしないようにすることが可能になっています。

公式リポジトリのdepth2imgプレビュー

1枚目が入力画像で、2枚目の白いやつが深度情報。3枚目以降が生成

上の画像は公式リポジトリにあるdepth2imgの例です。1枚目は入力画像で、MiDaSという深度推定モデルが入力画像のdepthを推測した結果が2枚目となります。この2枚の画像とテキストをUNetに入力することで、入力画像内にあるオブジェクト（ここではおじいさんですね）の位置関係が変わることも、背景の白い煙が何者かに変化することもなく、しかし画像内にあるオブジェクトに対して大きな変更を加えることができます。

.
.
.

しかし、ここで問題があります。
それは、こんなに優れたモデルなのにも関わらず、誰もこのモデルをファインチューニングしたものを公開していない（2023年1月13日現在）ため、計算資源や学習用データに恵まれない我々は、Stable Diffusion公式のdepthモデルが持つポテンシャルの範囲内でしか、深度情報を活用したスタイル変換モデルの実力を享受できないということです。

例えば上で行ったようなアニメ調への変換ですが、このdepthモデルはほとんどアニメ調の画像を学習していないのか、あまりよろしくない画像しか生成できません。

そして、十分なクオリティーでのアニメ調への変換を可能にするには、大量の学習用画像と潤沢な計算資源を集め、高い電気代を払ってこのモデルに新たなドメインを学習させなければなりません。

先月、このモデルのDreamboothスクリプトが親切な誰かによって公開されたので、これを使えば一応学習は進められますが、depth2imgを用いてやりたいことが公式のdepthモデルの知っているドメインとかけ離れている場合、Dreamboothのような小規模の学習で公式モデルを目的のドメインに適応させるのはかなり難しく、現状はこのモデルを望みのドメインに対して大規模に追加学習してくれて、なおかつ無料で公開してくれる神待ちの状態と言っても過言ではありません。

という感じで、例に漏れず自分も大規模な学習を行えるような計算資源を持ち合わせていないため、なんとかして小手先だけで上手くやれないかと考えました。

Task Arithmetic について

ここで、この論文が頭をよぎります。

https://arxiv.org/abs/2212.04089

アブストラクトを引用します。

Changing how pre-trained models behave—e.g., improving their performance on a downstream task or mitigating biases learned during pre-training—is a common practice when developing machine learning systems. In this work, we propose a new paradigm for steering the behavior of neural networks, centered around task vectors.
A task vector specifies a direction in the weight space of a pre-trained model, such that movement in that direction improves performance on the task.
We build task vectors by subtracting the weights of a pre-trained model from the weights of the same model after fine-tuning on a task. We show that these task vectors can be modified and combined together through arithmetic operations such as negation and addition, and the behavior of the resulting model is steered accordingly. Negating a task vector decreases performance on the target task, with little change in model behavior on control tasks.
Moreover, adding task vectors together can improve performance on multiple tasks at once. Finally, when tasks are linked by an analogy relationship of the form “A is to B as C is to D”, combining task vectors from three of the tasks can improve performance on the fourth, even when no data from the fourth task is used for training. Overall, our experiments with several models, modalities and tasks show that task arithmetic is a simple, efficient and effective way of editing models.

ざっくり言うと、

事前学習済みモデルの重みとファインチューニングされたモデルの重みの差をタスクベクトルとする。
タスクベクトルは、事前学習済みモデルの重み空間において、そのモデルがあるタスクに対する性能を向上させる方向を定めるもの。
同一の事前学習済みモデルから追加学習されたモデル同士であれば、タスクAについて追加学習されたモデルから抽出したタスクベクトルを別モデルに加算することによって、別モデルのタスクAの性能が向上する。（逆も然り）
タスクベクトルを足し合わせることで、複数のタスクに対する性能を向上させることができる。

という内容です。タスク減算・類推の話はかなり示唆に富むものがあり、内容も結構軽いのでかなりおすすめの論文です。

話を戻します。
つまり、この論文から今回の話に繋がることとして、

同一の事前学習済みモデルからファインチューニングされた、それぞれ別々のモデルから抽出したタスクベクトルを足し合わせたモデルは、その別々のタスクに対する性能が向上する

ということになります。

ここまでの話を聞いていて、Stable Diffusionから追加学習されたイラスト生成モデルをよく触っている方々はなんとなく心当たりがあるかもしれませんが、これはモデルのマージとめちゃくちゃ似ている話です。

しかし二点違っていて、一つはTask Arithmeticがモデルとモデルの加重平均ではないことにあります。
そのためタスクベクトルは、その作成のベースとなったモデルが同一であれば、比率の合計が1にならないような重みづけで加算・減算できます。

またもう一点について、あるタスクに特化したモデルの構造が、追加学習元のベースモデルの構造と比較して少し変わったりしていても、双方で共有しているレイヤーの重み空間の形状が根本的に異なるわけではなければ、タスクベクトルは（おそらく）両モデルの共通部分において足し合わせることができるはずです。
一方、モデルマージは両者同じレイヤーを持っていないと行えません。

それでは、depthモデルのベースになったモデルを見返してみましょう。

This stable-diffusion-2-depth model is resumed from stable-diffusion-2-base (512-base-ema.ckpt) and finetuned for 200k steps. Added an extra input channel to process the (relative) depth prediction produced by MiDaS (dpt_hybrid) which is used as an additional conditioning.
https://huggingface.co/stabilityai/stable-diffusion-2-depth

新しい入力チャンネルが追加されていて、重み空間の形状がかなり変化している可能性があるものの、ベースモデルはStable Diffusion V2.0らしいです。
ということで、ここまでの考えを踏まえると、depthモデルはStable Diffusion V2.0から追加学習されたモデルのタスクベクトルと加算することができるのではないか？というのが今回の仮説で、試したことです。

ここまででもう大体分かったという方は、ぜひ手元で試してみてください。

別モデルのタスクベクトルをdepthモデルに加算する

まず、用意する別モデルですが、最低限Stable Diffusion V2系から追加学習されたモデルである必要がある（SD-V1系では加算できないことを確認済み）ため、Waifu Diffusion V1.4にします。

なお、Waifu Diffusion V1.4のタスクベクトルをdepthモデルに加算することは、Stable Diffusion V2.0をベースモデルとして、depthモデルのタスクベクトルとWaifu Diffusion V1.4のタスクベクトルを加算することと同義です。

タスクの加算を行うにあたってまず、本当に加算できるのか？を確認します。双方の重み空間の形状があまりにも似ていなかったらおそらく加算はできないということが分かるので、双方のUNetで形状が同じ各レイヤーのコサイン類似度を見てみましょう。

import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

depth = torch.load("./512-depth-ema.ckpt")
waifu1_4 = torch.load("./wd-1-4-anime_e2.ckpt")

sim_dict = {}

for key in set(depth["state_dict"].keys()) & set(waifu1_4["state_dict"].keys()):
    # skip weight of an input block whose dimension is different from each other
    if "model.diffusion_model" in key and "model.diffusion_model.input_blocks.0.0" not in key:
        depth_weight = depth["state_dict"][key].numpy().flatten()
        waifu1_4_weight = waifu1_4["state_dict"][key].numpy().flatten()
        # calc cosine similarity
        cos_sim = (
            np.dot(depth_weight, waifu1_4_weight) 
            / (np.linalg.norm(depth_weight) * np.linalg.norm(waifu1_4_weight))
        )
        sim_dict[key] = cos_sim

df = pd.DataFrame.from_dict(
    sim_dict, 
    orient="index", 
    columns=["cosine_similarity"],
)
    
df.plot(
    kind="hist", 
    figsize=(10, 10), 
    ylabel="count", 
    xlabel="cosine similarity", 
    title="Cosine similarity between depth and waifu1-4",
)
plt.show()

かなり似ていますね。いけんちゃうか？という感じなので、加算してみましょう。
なお、Waifu Diffusion V1.4のVAEはStable Diffusionのそれよりアニメ調に適しているため、そこだけは加算ではなく、置き換えます。

import copy
import torch

base = torch.load("./512-base-ema.ckpt")
depth = torch.load("./512-depth-ema.ckpt")
waifu1_4 = torch.load("./wd-1-4-anime_e2.ckpt")

# make task vector
task_waifu = copy.deepcopy(waifu1_4)
for key in set(waifu1_4["state_dict"].keys()):
    if "model.diffusion_model" in key:
        task_waifu["state_dict"][key] = waifu1_4["state_dict"][key] - base["state_dict"][key]
    
for key in set(depth["state_dict"].keys()) & set(task_waifu["state_dict"].keys()):
    # replace weight of VAE with waifu's weight
    if "first_stage_model" in key:
        depth["state_dict"][key] = task_waifu["state_dict"][key]
    # don't replace weight of an input block whose dimension is different from each other
    elif "model.diffusion_model.input_blocks.0.0" in key:
        pass
    # add task weight to unet layers
    elif "model.diffusion_model" in key:
        task_depth = depth["state_dict"][key] - base["state_dict"][key]
        depth["state_dict"][key] = base["state_dict"][key] + (task_depth * 0.5 + task_waifu["state_dict"][key] * 0.5)
	
torch.save(depth, "./depth_waifu_1_4.ckpt")

それでは、早速このモデルでdepth2imgをやってみます。
入力は上のものと同じ画像・テキストを使ってみましょう。

(再掲)入力画像

以下がdepth2imgの結果です。

変換後

どうでしょうか？実写→アニメ調への変換はバッチリですし、画像内のオブジェクトが何か全く別のモノと解釈されて置き換わっているということもありません！（人物が似てないのは入力テキストのほうで似せる努力をしていないため）

特に大掛かりなことをするでもなく、depthモデルとWaifu Diffusion V1.4、それぞれが得意なタスクを両方ともこなせるモデルを生み出すことができました。

ちなみに、このキメラモデルを応用してこんなこともできます（1枚目が入力、2枚目が出力）。

なお、加算のときの係数である0.5, 0.5には何の根拠もないため、気になる方がいたら色々数値を変えて試してみてください。

.
.
.

おわりに

depthモデルはStable Diffusion V2.0のファインチューニングモデルですが、オリジナルから構造が少し変わっているので、まさかTask Arithmeticが使えるとは正直思っていなくて、ダメ元で適当にやったら上手くいってビビり散らかしました。

当然ですが、Waifu Diffusion V1.4に限らずStable Diffusion V2.0から追加で学習されたモデルであれば、同様の手法で組み合わせることが出来ると思います。楽しいですね。

ぜひ試してみてください。悪用は🆖

はじめに

Depth-to-Imageモデルの説明

Task Arithmetic について

別モデルのタスクベクトルをdepthモデルに加算する

おわりに

Discussion