【論文読み】フレーム間の一貫性を保持した人物動画生成モデル: Animate Anyone

この記事は[京都大学人工知能研究会KaiRA Advent Calendar 2023](https://qiita.com/advent-calendar/2023/kyoto-kaira) 4日目の記事になります。

先日公開された動画生成モデル「Animate Anyone」の論文についてまとめました。

arXiv

https://arxiv.org/abs/2311.17117

プロジェクトページ

https://humanaigc.github.io/animate-anyone/

デモ動画

https://www.youtube.com/watch?v=8PCn5hLKNu4

Animate Anyoneでは1枚の人物画像とポーズ動画を入力に受け取り、そのポーズ動画の通り人物画像を動かすということをやっています。
デモ動画を見てもらうとわかる通り、既存手法に比べてフレーム間のちらつきが格段に抑制されています。TikTokデータセットでは若干微妙ですが、入力が綺麗な画像の場合には実写動画レベルの品質で生成出来ています。また、アニメ画像や人間でない画像についても同様に動きを付けることが可能になっています。

## 手法概要
Animate Anyoneの重要なコンポーネントとして以下の3つが挙げられています。
 - 入力画像から情報を抽出するための**ReferenceNet**
 - ポーズ情報を抽出するための**Pose Guider**
 - 時間方向の処理を行うための**Temporal Layer**

以下が論文から引用した概要図です。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/629111/7ef7d1b1-7573-2a47-5989-8bc60d445cbc.png)

本手法ではStable Diffusion (SD)をベースとしています。
SDはDiffusion Modelをピクセル空間ではなく潜在空間で行っており、ピクセル空間と潜在空間の変換は別途学習されたEncoderとDecoder(VAE)によって行われます。Diffusion Modelのモデル構成ではUNetを採用しており、Convolution、self-attentention、cross-attention等のモジュールによって構成されます。cross-attention部分ではCLIPのText Encoderによって計算されたテキスト埋め込みと特徴マップとのattentionが計算されることになります。

### ReferenceNet
参照画像の情報を反映するためのモジュールです。
今回は参照画像の情報を反映して動画を生成する必要があり、テキストで条件づけるだけでは不十分だと考えられます。既存手法ではCLIPのText Encoderに加えてCLIPのImage Encoderを使って参照画像の埋め込みも入力していました。しかし、
 - 224×224にリサイズされて処理されるため画像の詳細が潰れてしまう
 - そもそもCLIPはテキストとの対応を学習させているため、テキストで表せないような細かい情報が反映されない

という理由から、CLIPのImage Encoderだけでは参照画像の情報を十分に反映できません。

そこで、ReferenceNetを用いて、self-attentionによって参照画像の情報を詳細に反映することを考えます。ReferenceNetとAnimate Anyoneのデノイジングを行う方のUNetの構造は(Temporal Layerを除くと)同じで、どちらもSDのUNetを初期値としています。ReferenceNetでは参照画像が入力された時のself-attentionへの入力を保持しておき、デノイジングを行う方のUNetではself-attentionの計算時にReferenceNetの方で保持していた入力と一緒にattentionを計算します。
PyTorch-likeに書けば、ReferenceNetの方のself-attentionの入力を`x1`、UNetの方のself-attentionの入力を`x2`とすると、それぞれのshapeは`(N,C,H,W)`、`(N,T,C,H,W)`となり(`T`はフレーム数です)、`torch.cat([x1.repeat(T,1,1,1), x2.reshape(N*T,C,H,W)], dim=3)`としてspatialなself-attentionが計算されます。
また、cross-attentionの方ではImage Encoderの出力とのattentionを計算しています。

なお、SDにテキスト以外の条件を追加する方法として真っ先に候補に挙がるのがControlNet(特にinpaintのControlNet)ですが、ControlNetは条件画像と生成したい画像とが空間的に整合性が取れている必要があるため適していないとのことです。(確かにControlNetは特徴量マップを元のSDに足し合わせることで条件を追加する構成になっているため、今回のように参照画像と生成したい画像でポーズが異なっている場合には相性が悪そうです。) Ablation Studyで実際にReferenceNetの代わりにControlNetを使用した場合を実験していますが、服の柄が変わるなどして上手く行っていません。

### Pose Guider
ポーズ情報をUNetに反映するためのモジュールです。
こちらも真っ先に候補に挙がる方法はControlNetですが、ControlNetはデノイジングの過程で毎ステップ計算が必要になるため処理が重いという欠点があります。その代わりに、Pose Guiderという軽量なネットワークを用意し、その出力をUNetへの入力に足し合わせるという方法を取っています。また、ControlNet同様出力層のConvolutionの重みを0で初期化しており、学習初期は元のSD同様ポーズ情報が含まれない入力がUNetに渡されることになります。

### Temporal Layer
フレーム間の一貫性を保持するために、フレーム方向にattentionを計算するモジュールです。デノイジングする方のUNetのself-attentionとcross-attentionの後に追加されます。先ほどと同様にPyTorch-likeに書けば、shapeが`(N,T,C,H,W)`の`x`に対して、`x.permute(0,3,4,1,2).reshape(N*H*W,T,C)`として`T`の次元に沿ってself-attentionを計算します。
また、このモジュールはAnimateDiffで使われているものと同じであり、初期値としてAnimateDiffの方で学習させた重みを使っています。

### 学習
学習は2段階に分けて行われます。
1段階目は、Temporal Layerがない状態で学習させます。データセット中の動画から2フレーム取り出し、片方をRefereceNet、もう片方からポーズ情報を取り出してUNetの方に入力します。パラメータの更新はRefereceNetとUNetのみに対して行われます。この学習により1枚の画像を与えられたポーズの画像に変換することが可能になります。
2段階目は、Temporal Layerを追加して学習させます。UNetの方には24フレーム分をconcatしたものを入力します。この段階ではTemporal Layerのパラメータのみを更新します。この学習によりフレーム間の一貫性が向上します。

推論時には、DDIMを用いて20ステップで動画を生成します。

### 実験
定性評価としては論文を見ての通り、提案手法の質が最も高いと言えます。
定量評価では、SSIM、PSNR(SN比)、LPIPS、FVDという評価指標で比較されており、すべての指標で比較手法よりも良い結果を出しています。特に動画レベルの評価指標であるFVD(前者3つは静止画レベルでの評価です)の改善が著しいです。

また、Ablation StudyではRefereceNetの代わりにCLIPとControlNetを用いた場合を実験していますが、明らかに服の柄が変わっていたり、参照画像の情報を正確に反映できていません。

### Limitations
Limitationとして3点あげられています。
 - 他のモデル同様、手の生成品質が悪い場合がある
 - 参照画像が1つしか与えられないため、写っていない部分の生成が不良設定問題である
 - デノイジングの反復的な計算に時間がかかる

## 感想
 - 学習コストがどれくらいなのか気になりました。学習済みのSDがベースなので、スクラッチから学習させる手法よりは時間がかからなさそうですが…
 - 例えば参照画像を追加してattention maskを上手いこと設定してあげれば、首から上はこの画像、服はこの画像、というように部分的に指定することも可能なのでは
 - 最近SDを1-stepやそこらで生成する手法はいくつか出てきているので、それをベースに学習した場合どうなるのか気になります。