【自動タイムラプス製造!?】Paints-UNDO技術解説

毎度お馴染みlllyasvielさんがまた技術革新を起こしているので、最速解説目指して記事を書いていきます。
今回解説する技術はこちら
https://github.com/lllyasviel/Paints-UNDO

こちらからデモが見れます
https://lllyasviel.github.io/pages/paints_undo/

ここではどのような技術か？を重点的に解説していこうと思います。

Paints-UNDOを支える二つのモデル

single-frame model

single-frame modelは、1つの画像とoperation stepという数字情報を入力として受け取り、1つの画像を出力するモデルになります。
これは、一つのイラストが完成するまでに1000回人間が操作 （ここでいう操作とは、ブラシストロークなどを指しているらしいです。もっというと、Ctrl-Zで戻る差分を操作とみなしてよさそう） を行うという仮定を置き、
operation stepが999（何も描かれていない真っ白なキャンバスに最初に書き込まれたストローク）から始まり、operation stepが減る毎に操作が一回分追加され、operation stepが0、即ち1000回目の操作が行われた時に絵が完成する、という動きになっているようです。

multi-frame model

multi-frame modelは、2つの画像を入力として受け取り、入力された画像間の16個の中間フレームを出力するモデルです。
いわゆる中割り生成モデルに近いイメージですが、生成される画像は動きを補完するフレームではなく、入力された一つ目の画像をスタートとし、二つ目の画像に辿り着くまでの作業過程を出力するモデルのようです。

single-frame modelと比較すると生成される画像の一貫性が非常に高いですが、生成速度が遅く、創造性も低く、生成できる画像も16個（16フレーム）に限定されてしまうというデメリットを持っています。

モデルの併用

Paints-UNDOはこの二つのモデルを組み合わせて構成されているようです。
併用の仕方としては、まずsingle-frame modelを使って5~7個のキーフレームを生成し、次にmulti-frame modelを使ってキーフレーム間の作業を補間する、といった処理を行っています。

各モデルの詳細解説

大枠の仕組みはわかったところで、より詳細なモデルの構成に触れていきます。
Paints-UNDOは論文などが公表されているわけではないので、ここから先は私がGithub上で公開されているコードを読みながら、おそらくこういう処理を行っているのだろう、と予想した内容になります。
そのため、正確性が乏しい可能性があるためご承知ください。

では、まずはsingle-frame modelから深掘っていきましょう

single-frame model

single-frame modelはStableDiffusion v1.5をベースとして作られています。
SD1.5との相違は、clipとbeta schedulerの変更、及び先ほど説明したoperation stepという生成条件が追加された点です。

clipはViT-L/14をベースとし、最後のレイヤー(12層目)を完全に削除しているようです。
なので、Diffusersを用いる際は、CLIP Skipを常に2にするように指定されています。

続いてbeta scheduler。
こちらは、元のSD1.5ではbetaは以下のように定義されています。

betas = torch.linspace(0.00085, 0.020, 1000, dtype=torch.float64)

一方、single-frame modelでは新たにこのように設定されています。

betas = torch.linspace(0.00085 ** 0.5, 0.012 ** 0.5, 1000, dtype=torch.float64) ** 2

二つの定義を見比べると乗算が行われている点が大きく異なります。
なぜこのようにしたかは、The choice of this scheduler is based on our internal user study.とのことなので、この定義だと経験的に上手くいく、程度の理解で良いと思います。

最後に、肝となるoperation stepについて。
こちらは、SDXLのExtra Embeddingと同様の方法で、埋込み層に生成条件を追加しているとのこと。
これだけしか説明がないので、多分Textual Inversionとかと同じイメージなのかな？と思います。
が、あくまで推測なので実態を確認するためにoperation stepを読み込んでいる部分のコードを実際に見てみましょう。

まずは、モデルの呼び出し部分です。(single-frame modelに関わる部分のみ抜粋)

    #入力画像
    fg = resize_and_center_crop(input_fg, image_width, image_height)
    concat_conds = numpy2pytorch([fg]).to(device=vae.device, dtype=vae.dtype)
    concat_conds = vae.encode(concat_conds).latent_dist.mode() * vae.config.scaling_factor
    concat_conds = concat_conds.to(device=unet.device, dtype=unet.dtype)

    #プロンプト
    conds = encode_cropped_prompt_77tokens(prompt)
    unconds = encode_cropped_prompt_77tokens(n_prompt)

    #Operation Step
    fs = torch.tensor(input_undo_steps).to(device=unet.device, dtype=torch.long)
    initial_latents = torch.zeros_like(concat_conds)
    

    #潜在空間生成
    latents = k_sampler(
        initial_latent=initial_latents,
        strength=1.0,
        num_inference_steps=steps,
        guidance_scale=cfg,
        batch_size=len(input_undo_steps),
        generator=rng,
        prompt_embeds=conds,
        negative_prompt_embeds=unconds,
        cross_attention_kwargs={'concat_conds': concat_conds, 'coded_conds': fs},
        same_noise_in_batch=True,
        progress_tqdm=functools.partial(progress.tqdm, desc='Generating Key Frames')
    ).to(vae.dtype) / vae.config.scaling_factor

ここで、input_undo_stepsが問題のoperation stepに該当する部分ですね。
input_undo_stepsはテンソル化されてfsという変数に格納されているため、追ってみるとcross_attention_kwargsという引数でSamplerに与えられています。
ここで、cross_attention_kwargsはconcat_condsとcoded_condsという二つのキーを持っていますが、これが何を意味しているのかを確認するために、モデルの定義部分を見てみましょう

class ModifiedUNet(UNet2DConditionModel):
    @classmethod
    def from_config(cls, *args, **kwargs):
        m = super().from_config(*args, **kwargs)
        unet_add_concat_conds(unet=m, new_channels=4)
        unet_add_coded_conds(unet=m, added_number_count=1)
        return m


model_name = 'lllyasviel/paints_undo_single_frame'
unet = ModifiedUNet.from_pretrained(model_name, subfolder="unet").to(torch.float16)

k_sampler = KDiffusionSampler(
    unet=unet,
    timesteps=1000,
    linear_start=0.00085,
    linear_end=0.020,
    linear=True
)

uentの定義部分でModifiedUNetというクラスが呼び出されています。
これはDiffusersから提供されているUNet2DConditionModelを拡張したもので、concat_condsとcoded_condsという画像生成時に生成条件を追加するための構造が追加されています。

では、各々の細かい実装を見てみましょう。
まずは、concat_condsについて。
concat_condsは以下のように定義されます。

def unet_add_concat_conds(unet, new_channels=4):
    with torch.no_grad():
        new_conv_in = torch.nn.Conv2d(4 + new_channels, unet.conv_in.out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding)
        new_conv_in.weight.zero_()
        new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight)
        new_conv_in.bias = unet.conv_in.bias
        unet.conv_in = new_conv_in

    unet_original_forward = unet.forward

    def hooked_unet_forward(sample, timestep, encoder_hidden_states, **kwargs):
        cross_attention_kwargs = {k: v for k, v in kwargs['cross_attention_kwargs'].items()}
        c_concat = cross_attention_kwargs.pop('concat_conds')
        kwargs['cross_attention_kwargs'] = cross_attention_kwargs

        c_concat = torch.cat([c_concat] * (sample.shape[0] // c_concat.shape[0]), dim=0).to(sample)
        new_sample = torch.cat([sample, c_concat], dim=1)
        return unet_original_forward(new_sample, timestep, encoder_hidden_states, **kwargs)

    unet.forward = hooked_unet_forward
    return

unet.conv_inが新しく定義されたnew_conv_inで上書きされ、入力できるChannelが追加されているのがわかります。
diffusersのunet.conv_inは元々以下のように定義されています。

self.conv_in = nn.Conv2d(
            in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding
        )

in_channelsはもとも Batch, height, width, RGBの4チャンネルであることを踏まえると、4 + new_channelsはデノイズ途中の画像に加えて、Conditioning画像(即ち、制作過程を作成したい完成画像)を受け入れ可能にするための構造とみなしてよさそうです。

c_concat = torch.cat([c_concat] * (sample.shape[0] // c_concat.shape[0]), dim=0).to(sample)
new_sample = torch.cat([sample, c_concat], dim=1)
return unet_original_forward(new_sample, timestep, encoder_hidden_states, **kwargs)

実際、上記のコードを見てみると、sample(デノイズ途中の画像)に、c_concat(生成条件となる過程を作成したい完成画像)をcatで繋げたあと、通常のforwardに投げているのがわかります。

続いてcoded_condsです。
こちらがoperation stepを生成条件に追加する部分ですね
unet_add_coded_condsの定義は以下のようになります。

def unet_add_coded_conds(unet, added_number_count=1):
    unet.add_time_proj = Timesteps(256, True, 0)
    unet.add_embedding = TimestepEmbedding(256 * added_number_count, 1280)

    def get_aug_embed(emb, encoder_hidden_states, added_cond_kwargs):
        coded_conds = added_cond_kwargs.get("coded_conds")
        batch_size = coded_conds.shape[0]
        time_embeds = unet.add_time_proj(coded_conds.flatten())
        time_embeds = time_embeds.reshape((batch_size, -1))
        time_embeds = time_embeds.to(emb)
        aug_emb = unet.add_embedding(time_embeds)
        return aug_emb

    unet.get_aug_embed = get_aug_embed

    unet_original_forward = unet.forward

    def hooked_unet_forward(sample, timestep, encoder_hidden_states, **kwargs):
        cross_attention_kwargs = {k: v for k, v in kwargs['cross_attention_kwargs'].items()}
        coded_conds = cross_attention_kwargs.pop('coded_conds')
        kwargs['cross_attention_kwargs'] = cross_attention_kwargs

        coded_conds = torch.cat([coded_conds] * (sample.shape[0] // coded_conds.shape[0]), dim=0).to(sample.device)
        kwargs['added_cond_kwargs'] = dict(coded_conds=coded_conds)
        return unet_original_forward(sample, timestep, encoder_hidden_states, **kwargs)

    unet.forward = hooked_unet_forward

    return

細かく見ていきましょう。
まず、以下の部分でUNetに時間方向の埋め込み層を追加しています。

unet.add_time_proj = Timesteps(256, True, 0)
unet.add_embedding = TimestepEmbedding(256 * added_number_count, 1280)

TimeStepEmbeddingは、Diffusion Modelの逆拡散過程を制御する際にも使われており、今どのステップか？をモデルに理解させる時に使われるものになります。

続いてget_aug_embed。

def get_aug_embed(emb, encoder_hidden_states, added_cond_kwargs):
    coded_conds = added_cond_kwargs.get("coded_conds")
    batch_size = coded_conds.shape[0]
    time_embeds = unet.add_time_proj(coded_conds.flatten())
    time_embeds = time_embeds.reshape((batch_size, -1))
    time_embeds = time_embeds.to(emb)
    aug_emb = unet.add_embedding(time_embeds)
    return aug_emb

unet.get_aug_embed = get_aug_embed

ここで、coded_condsとして渡された「完成された絵ができる過程のうち、どのステップを生成したいか？」という条件を追加しています。
get_aug_embedは元々DiffusersのUNet2DConditionModelが持っている関数で、そこを新しく定義したget_aug_embedでオーバーライドしている形になります。

get_aug_embed自体は、通常のforward関数内で呼び出されるため、unet_add_coded_conds内で呼び出されることはありません。

kwargs['added_cond_kwargs'] = dict(coded_conds=coded_conds)
        return unet_original_forward(sample, timestep, encoder_hidden_states, **kwargs)

ここでkwargsとしてoperation stepをcoded_condsとして登録することで、get_aug_embedがforward関数内で呼び出された際に、生成したいoperation stepを渡すことができる、という仕組みになっています。

というような形で、制作過程を生成したい完成絵と、どのステップの過程を生成したいか？という条件になるoperation stepを、画像の生成条件として与えられるようになっていることがわかりました。
（正直、こんな簡単なモデル改造で実現できるとは素直に驚きました）

multi-frame model

キーフレームを生成するSingle-flame modelの解説は以上にして、キーフレーム間を補間する動画生成モデルについて解説をしていきます。
ただし、multi-frame modelのコードベースの解説はToonCrafterを理解していないといけない + 動画生成モデルとしての精度向上がメインで本筋とはあまり関係ないので、コードには触れずにさらっと解説をします。
（ちなみに、ToonCrafterの進化版として普通に使えそうなので、アニメーションのフレーム補間などにもmulti-frame modelとしても使える気がします。要実験）

multi-frame modelは、最近話題になったToonCrafterの大元になっている、VideoCrafterのモデルをベースとして用いているようです。
ちなみに、VAEはアニメ生成用に調整されたVideoCrafterを用いているとのこと。
（閑話ですが、私は動画生成はDiTベースが主流になるだろうからVideoCrafter系列は追わなくて良いやろ！という舐めた考えを持っていました...。今回の技術が出てきたことで、VideoCrafter系列もやっぱり追わなきゃダメだなと反省しています...。どう頑張っても時間が足りない）

なお、ベースといっても初期の重みをVideoCrafterのモデルから流用しているだけで、ニューラルネットワークの構造は元のものからは大幅に変更しているとのこと。
トレーニングコードも推論コードも完全にゼロから実装し直したとのことで、lllyasvielさんの凄さがここでもわかりますね本当に同じ人間か？

全体的な構造としては、ToonCrafterの構造をベースに、3D-UNet、VAE、CLIP、CLIP-Vision、Image Projection の5つのコンポーネントを追加/変更したような形になっています。

VAEは前述したようにVideoCrafterが作成したアニメ生成に特化したモデルを使用しています。

3D-Unetもいくつか変更がされているようです。
大きな点では、Unetを訓練するようにした点と、Spatial Self Attention layersのtemporal Windowをサポートするようになった点です（わからん）
この辺りはToonCrafterの構造を一度勉強した後に再度追記できればと思いますが、おそらくSpatial Self Attention(即ち、空間に対する注意)に対し、時間的な一貫性を付与するような構造だと思います。
すごくざっくりいうとちらつきとか時間経過による破綻の抑制に寄与する部分なのかなと（違ったら後で追記します）

ちなみに、コード内のdiffusers_vdm.attention.CrossAttention.temporal_window_for_spatial_self_attentionとtemporal_window_typeをいじることで、以下の3モードを推論時に使い分けることができるようです（有効にするとGPUメモリの消費量が上がるため、デフォルトでは全部オフになっている）

• prv mode
  • 前のフレームとの空間的な一貫性を重視して画像(フレーム)を生成する
• first mode
  • 一番最初のフレームとの空間的な一貫性を重視して画像(フレーム)を生成する
• roll　mode
  • 前のフレームと次のフレームの空間的な一貫性を重視して画像(フレーム)を生成する

ぱっと見roll modeが一番品質の良い動画を生成できそうですが、その分VRAM消費量も上がりそうなので、どのモードを使うかはお手持ちの計算リソースと要相談ですね

続いてCLIP。
CLIPはSD2.1のCLIPを使っています、以上。

CLIP繋がりで、CLIP-Visionという技術も使われています。
こちらは、 位置埋め込みを補間することで任意のアスペクト比をサポートする Clip Vision (ViT/H) とのことです。
要はどんなアスペクト比の画像でもCLIPで認識できるようにする技術かなと。
CLIPが受け入れられる形式が224x224なので、通常は入力画像を224x224にリサイズするか、ToonCrafterのように中心部を224x224で切り抜くかするのですが、今回はそういった手法は使っていないよとのこと。

最後にImage Projection。
これは2つの画像を入力として受け取り、その画像の中間フレームを16枚生成できるよう、小型のTransformerを新しく実装したとのこと。
普通にこれだけでアニメーション補間の新しい選択肢になるので、タイムラプス生成関係なく中割り生成モデルとして使ってみたいですね。

終わりに

ということで、今回は新しく発表された自動タイムラプス生成ツールことPaints-UNDOの技術解説でした！

タイムラプス生成はおまけで、本イラストを生成する過程を1000ステップに分解し、各ステップごとに生成可能にした画像生成モデルの発明と、既存モデルよりも高品質な中割り生成モデルの開発というのが本質だったのかなと思います。

single-frame modelもmulti frameモデルも、それ単体だけで色々と新しいことができそうなので、またこの技術をベースとして画像/動画生成技術もより飛躍していくのでしょう。

ここまで読んでいただきありがとうございました！