StoryDiffusion Consistent Self-attention for Long-range Image And Video Generation

StoryDiffusion: Consistent Self-attention for Long-range Image And Video Generation

TL; DR：一阶段通过生图 batch 内自注意 KV 共享，生成主题一致的图片。二阶段根据这些图片进行特征插值并作为条件生成中间帧，最后得到视频。

---

方法

Story Diffusion 整体上可以分为两个阶段，第一阶段在生图模型中插入 Consistent Self-Attention（CSA）模块，来生成主题一致的图片。这一阶段是不需要训练的。第一阶段生成的主题一致图片就直接可以用来拼成故事画了，也可以输入到第二阶段。第二阶段，根据第一阶段生成的主体一致图片，生成一致转场视频。

第一阶段

一个故事中有多个故事分镜，每个文本描述 prompt 对应一个分镜图片。Story Diffusion 第一阶段要做的事情就是在生成多张背景不同的分镜图片时，保持故事中的（多个）人物在所有分镜图片中的形象（脸部、服装、配饰等）保持一致。

Story Diffusion 第一阶段的 CSA 与 ConsiStory 中的 SDSA 思路类似，即通过在同一个 batch 内，UNet 的注意力层在各图片之间进行 KV 共享，使得 batch 内不同图片可以彼此交互，从而实现主题一致的生成结果。但是 CSA 取交互 KV 的方式与 SDSA 有所不同，是在同 batch 内随机找另一张图片取 KV。而且 CSA 的 KV 共享只做在 UNet 的自注意力层中，交叉注意力层不做。

第一阶段，生成主题一致图片的核心模块 CSA 及其如何插入到标准 UNet 中的示意图如下所示。具体来说，记一个 batch 内的图片特征为 $\mathcal{I}\in {\mathbb{R}}^{B\times N\times C}$ ，其中 $B,N,C$ 分别为 batch size，每张图片的 token 个数以及通道数。标准的自注意力是在每张图片的特征自身进行的，表示为：
$$O_i=\text{Attention}(Q_i,K_i,Vi)$$
其中 $Q_i,K_i,V_i$ 分别是第 $i$ 个图片特征 $I_i$ 经过线性变换得到的。

在 Story Diffusion 的 CSA 中，为了实现 batch 内的特征交互，对第 $i$ 个图片特征，首先从 batch 内的其他图片特征中随机选一个：
$$S_i=\text{RandSample}(I_1,\dots ,I_{i-1},I_{i+1},\dots ,I_B)$$
然后将 $S_i$ 与原特征 $I_i$ 拼接起来得到 $P_i$ ，再经过线性映射，得到新的 key 和 value：$K_{P_i},V_{P_i}$ 。再进行注意力计算：
$$O_i=\text{Attention}(Q_i,K_{P_i},V_{P_i})$$
这样，batch 内的图片在生成时能够相互交互，从而保持主题的一致。

另外，为了解决在生成较长故事时 batch 过大，显存不足的问题，Story Diffusion 还采用了滑窗机制来进行生成，具体伪代码如下所示：

第二阶段

Story Diffusion 的第二阶段是在第一阶段生成的一系列主题一致图片的基础上，进一步生成视频。也就是在每两张图片之间插帧，生成连贯的视频，相当于是一个给定首尾帧，生成中间帧的视频生成任务。作者指出现有的工作（SparseCtrl 和 SEINE），在首尾两帧差异过大时，生成效果不够稳定。作者认为这是因为他们只依赖于时序模块来生成中间帧，所以对首尾帧差距太大的情况处理得不够好。

本文提出的视频插帧方法如下图所示。对于给定的首尾两帧图片，首先通过一个图像编码器（CLIP）提取它们的特征。然后通过线性插值，得到中间帧的特征。并经过 n 层 Transformer Block 进行处理。再将各帧的特征作为条件，（与 prompt 文本 embeddings 拼接起来）通过交叉注意力的形式注入到 UNet 的生图过程中。

总结

Story Diffusion 第一阶段的 batch 内 KV 共享来实现图片彼此之间的交互，进而保证主题一致性。第二阶段先提取首尾帧特征并插值，处理后作为条件注入到生图过程中。

第一阶段的 KV 共享的方式感觉不如 ConsiStory 的方式精妙，但是 ConsiStory 是为了背景的多样性选择了仅在主题区域内进行 KV 共享，Story Diffusion 的随机选图全局共享的方式，背景一致性也得到了一定保证，对后续的插帧视频生成更友好一些。还有个问题就是第二阶段为啥要用 CLIP 语义特征，以图 3 为例，首尾两帧的语义特征应该是非常接近的（比如 “一个人走在马路上”），似乎用纹理特征更合理一些？