Skip to main content

Stable Video Diffusion の公式実装を読む

· 9 min read

Stable Video Diffusion の推論処理を、Stability AI の公式実装である Stability-AI/generative-models をもとに解説していく。

全体像​

推論スクリプトと YAML config を組み合わせて、DiffusionEngine の中に部品を組み立てていく構成になっている。 メインの実装は simple_video_sample.py の sample() にある。ここで入力画像の読み込みから出力動画の保存までを行っている。

version から config を選ぶ​

sample() の前半では、versionと呼ばれるパラメータに応じて

  • フレーム数 (出力動画の長さ)
  • ステップ数 (デノイジングの回数)
  • config ファイル

を選択する。公式実装では、svd と svd_xt の2種類が用意されている。

svd     -> 14 frames, 25 steps, svd.yaml
svd_xt  -> 25 frames, 30 steps, svd_xt.yaml

その後、load_model() が YAML を読み、instantiate_from_config() でモデルを作る。svd.yaml のルートを見ると、中心は DiffusionEngine で、その下に denoiser、network、conditioner、first stage、sampler などが存在する構成になっている。

DiffusionEngine
  ├─ Denoiser
  ├─ VideoUNet
  ├─ GeneralConditioner
  ├─ AutoencodingEngine / VideoDecoder
  └─ EulerEDMSampler

入力画像は2つの条件経路に分かれる​

SVD を読む上で一番大事なのは、入力画像が1回だけ使われるわけではない、という点だと思う。公式実装では、画像から大きく2種類の条件を作る。

入力画像
  ├─ ノイズなし画像 -> OpenCLIP -> cross-attention 条件
  └─ 少しノイズを足した画像 -> VAE encoder -> concat 条件

スクリプト側では value_dict に条件素材を詰める。中身を意味で書くとこうなる。

cond_frames_without_noise = 元画像
cond_frames               = 元画像 + cond_aug * noise
fps_id                    = fps 条件
motion_bucket_id          = motion 条件
cond_aug                  = 画像に足したノイズ量

cond_aug は、入力画像を VAE latent にするときのノイズ量であり、同時に条件 vector としてもモデルに渡される。つまり「どれくらい劣化した条件画像から復元する設定なのか」もモデルに知らせている。

GeneralConditioner が条件を仕分ける​

svd.yaml の conditioner_config を見ると、SVD の条件は3種類に分かれている。

cond_frames_without_noise
  -> FrozenOpenCLIPImagePredictionEmbedder
  -> crossattn

cond_frames
  -> VideoPredictionEmbedderWithEncoder
  -> concat

fps_id / motion_bucket_id / cond_aug
  -> ConcatTimestepEmbedderND
  -> vector

GeneralConditioner は、embedder の出力次元によって保存先を決める。2次元なら vector、3次元なら crossattn、4次元なら concat という分類になる。

ここでできる条件を diffusers 風に見ると、crossattn は image_embeddings、concat は image_latents、vector は added_time_ids に近い。

crossattn: CLIP image embedding
concat:   VAE image latent
vector:   fps + motion_bucket_id + cond_aug

この時点で、入力画像は「意味的な画像 embedding」と「空間的な latent」の両方として使われる。SVD の image-to-video らしさは、かなりここに詰まっている。

CFG 用に条件あり・なしを作る​

次に get_unconditional_conditioning() で、conditional な c と unconditional な uc を作る。

公式実装では、unconditional 側で cond_frames と cond_frames_without_noise をゼロにする。つまり CFG では、画像条件ありの予測と、画像条件を消した予測を比較する。

uc: 画像条件をゼロにした条件
c:  画像条件を持った条件

その後、crossattn と concat はフレーム数ぶん repeat される。公式実装では内部表現が [B*T, ...] に寄っているので、1枚の条件画像を T フレームぶんに複製してから、batch と time をまとめた形に並べ替える。

diffusers だと [B, T, C, H, W] で見えることが多いが、公式実装では [B*T, C, H, W] として扱う場面が多い。この layout の違いを頭に入れておくと読みやすい。

生成対象は video latent のランダムノイズ​

条件ができたら、生成対象の video latent をランダムノイズから始める。

samples_z の初期値: [T, 4, H/8, W/8] の noise

ここで T はフレーム数、4 は latent channel、H/8 と W/8 は VAE の空間圧縮後の解像度である。

重要なのは、入力画像の latent をそのまま動画に伸ばしているわけではないこと。最終的に denoise される主役は、T フレームぶんのランダムな video latent である。入力画像はあくまで条件として横から入る。

sampler は denoiser 関数を受け取る​

公式スクリプトでは、sampler に model をそのまま渡すのではなく、denoiser 関数を作って渡す。

EulerEDMSampler
  -> Denoiser
  -> VideoUNet

additional_model_inputs として image_only_indicator と num_video_frames も渡される。後者は VideoUNet 内で、batch と time の軸を戻すために必要になる。

サンプリング本体は EulerEDMSampler で、sigma schedule に沿ってノイズを少しずつ落としていく。diffusers の感覚でいえば、for t in timesteps の loop で unet と scheduler.step() を回しているところに近い。

image latent は UNet 入力 channel に concat される​

SVD の読みどころはここだと思う。concat 条件、つまり VAE encode された入力画像 latent は、UNet の入力 channel として結合される。

OpenAIWrapper の役割を意味で書くとこうなる。

noisy video latent: [B*T, 4, H/8, W/8]
image latent:       [B*T, 4, H/8, W/8]
------------------------------------------------
VideoUNet input:    [B*T, 8, H/8, W/8]

実際、svd.yaml では VideoUNet の in_channels が 8、out_channels が 4 になっている。UNet は「今の noisy video latent」と「条件画像の latent」を見ながら、次に取り除くべきノイズを4 channel で予測する。

つまり入力画像は、CLIP embedding として attention に入るだけではない。latent 空間でも各フレームに並べて渡される。この2経路の条件付けが、SVD の image-to-video の核になっている。

CFG はフレーム方向に scale される​

CFG は LinearPredictionGuider が担当する。

式としてはよく見る形で、

pred = uncond + scale * (cond - uncond)

である。ただし公式実装では min_scale と max_scale を持ち、フレーム方向に scale を変化させる。svd.yaml では min_scale: 1.0、max_scale: 2.5 になっている。

最初から最後まで同じ guidance scale をかけるというより、動画のフレーム軸を意識した CFG になっているのが面白い。

VideoUNet は T フレームをまとめて denoise する​

VideoUNet は、見た目は2D latent を [B*T, C, H, W] として受け取るが、内部では num_video_frames を使って time 軸を復元する。

VideoResBlock は一度 [B, C, T, H, W] のような形に戻して、時間方向の畳み込みを入れる。SpatialVideoTransformer も、空間方向の処理に加えて temporal attention を持っている。

このため、SVD は1フレームずつ逐次的に作っているわけではない。

T フレームぶんの latent をまとめて denoise する

という理解のほうが近い。画像から「次の1枚」を予測し、それをまた入力して、という autoregressive な動画生成ではない。

decode で latent から動画フレームへ戻す​

sampler が返す samples_z はまだ latent なので、最後に first stage model で decode する。

samples_z: [T, 4, H/8, W/8]
  -> VideoDecoder
samples_x: [T, 3, H, W]

SVD の config では decoder に VideoDecoder が使われる。ここでも単なる画像 VAE decoder ではなく、時間方向を意識した decoder になっている。

その後、値域を [0, 1] に戻し、watermark と safety filter を通して、imageio で mp4 として保存する。diffusers 版だと pipeline は frames を返し、保存は export_to_video() など外側の helper に任せることが多い。

まとめ​

公式実装を1本の流れにすると、こう読める。

入力画像を [-1, 1] tensor にする
  -> value_dict を作る
  -> GeneralConditioner で条件を作る
     ├─ OpenCLIP image embedding -> crossattn
     ├─ VAE image latent -> concat
     └─ fps / motion / cond_aug -> vector
  -> CFG 用に c / uc を作る
  -> random video latent を作る
  -> EulerEDMSampler で denoise
  -> VideoUNet 入力で video latent と image latent を concat
  -> temporal block / temporal attention で T フレームをまとめて処理
  -> VideoDecoder でフレームに戻す
  -> mp4 に保存する

自分の理解として一番大事だったのは、SVD が入力画像を2つの経路で使っていることだった。

1. CLIP embedding として cross-attention に入れる
2. VAE latent として noisy video latent に concat する

そして生成される動画は、入力画像 latent の単純な延長ではなく、T フレームぶんのランダム video latent を denoise して得られる。入力画像は、その denoise の方向を強く縛る条件として効いている。

この見方を持ってから diffusers の StableVideoDiffusionPipeline.__call__() を読むと、_encode_image()、_encode_vae_image()、added_time_ids、prepare_latents()、unet()、scheduler.step() の対応がかなり見えやすくなる。