Welcome back Yanming!

And boy meets world

0%

Stable Video(3D) Diffusion

Posted on Edited on In

Tutorial

SVD 使用 EDM 的思想，将离散噪声改进成连续噪声。对于原来的 DDPM，输入的是 \(t\) 这个参数；而 EDM 论文指出，\(t\) 实际上表示的是噪声强度 \(\sigma_t\)，应该直接把 \(\sigma_t\) 输入进模型。

所以 EDM 的观点是：与其用离散的 \(t\) 训练一个只认识离散噪声强度的去噪模型，不如训练一个认识连续噪声强度 \(\sigma\) 的模型。这样在采样 \(n\) 步时，不再是选择离散的去噪时刻 [timestep[n], timestep[n-1], ..., 0]，而是可以选择连续的噪声强度 [sigma[n], sigma[n-1], ..., 0]。

Reference：yifanzhou 的 SVD 解析

Code

time_embed

1
2
3
4
5
6
time_embed_dim = model_channels * 4  # 1280
self.time_embed = nn.Sequential(
    linear(model_channels, time_embed_dim),
    nn.SiLU(),
    linear(time_embed_dim, time_embed_dim),
)

label_emb

1
2
3
4
5
6
7
8
9
self.label_emb = nn.Sequential(
    nn.Sequential(
        linear(adm_in_channels, time_embed_dim),
        nn.SiLU(),
        linear(time_embed_dim, time_embed_dim),
    )
)

emb = emb + self.label_emb(y)

U-Net Architecture

注意：下面图中标注的维度可能不正确，因为它们来自 2D Stable Diffusion 模型。

Conv_in

1
2
3
4
5
6
7
8
self.input_blocks = nn.ModuleList(
    [
        TimestepEmbedSequential(
            # create 2D convolution, conv_in
            conv_nd(dims, in_channels, model_channels, 3, padding=1)
        )
    ]
)

Attention / ResBlocks

VideoResBlock / TemporalSpatialResBlock（ResnetBlock3D）

2D spatial resnet

1
2
3
4
5
6
7
8
9
10
11
12
super().__init__(
    channels,
    emb_channels,
    dropout,
    out_channels=out_channels,
    use_conv=use_conv,
    use_scale_shift_norm=use_scale_shift_norm,
    dims=dims,
    use_checkpoint=use_checkpoint,
    up=up,
    down=down,
)

3D temporal resnet

1
2
3
4
5
6
7
8
9
10
11
12
13
14
self.time_stack = ResBlock(
    default(out_channels, channels),
    emb_channels,
    dropout=dropout,
    dims=3,
    out_channels=default(out_channels, channels),
    use_scale_shift_norm=False,
    use_conv=False,
    up=False,
    down=False,
    kernel_size=video_kernel_size,
    use_checkpoint=use_checkpoint,
    exchange_temb_dims=True,
)

alpha blender

1
2
3
4
5
self.time_mixer = AlphaBlender(  # mixer
    alpha=merge_factor,
    merge_strategy=merge_strategy,
    rearrange_pattern="b t -> b 1 t 1 1",
)

SpatialVideoTransformer（Transformer3DModel）

Original 2D Transformer（SpatialTransformer）

1
2
3
4
5
6
7
8
9
10
11
12
super().__init__(
    in_channels,
    n_heads,
    d_head,
    depth=depth,
    dropout=dropout,
    attn_type=attn_mode,
    use_checkpoint=checkpoint,
    context_dim=context_dim,
    use_linear=use_linear,
    disable_self_attn=disable_self_attn,
)

Preprocessing layer

这里用 Linear 替代了卷积：

1
2
3
4
5
6
7
if not use_linear:  # use_linear is True
    self.proj_in = nn.Conv2d(
        in_channels, inner_dim, kernel_size=1, stride=1, padding=0
    )
else:
    # 使用 Linear 层代替 Transformer2DModel 中的 Conv1x1
    self.proj_in = nn.Linear(in_channels, inner_dim)

BasicTransformer

这里的维度与下面的图不同，因为下图表示的是原始 SD 模型。

注意有一个参数 self.disable_self_attn：如果为 true，第一个 self-attention 层会变成 cross-attention，也就是两个 attention 层都 conditioned on CLIP encoder 的输出。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
self.transformer_blocks = nn.ModuleList(  # 正式的 BasicTransformerBlock
    [
        BasicTransformerBlock(
            inner_dim,
            n_heads,
            d_head,
            dropout=dropout,
            context_dim=context_dim[d],
            disable_self_attn=disable_self_attn,
            attn_mode=attn_type,
            checkpoint=use_checkpoint,
            sdp_backend=sdp_backend,
        )
        for d in range(depth)
    ]
)

Post-process layer

这里同样用 Linear 替代卷积。

1
2
3
4
5
6
if not use_linear:
    self.proj_out = zero_module(
        nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
    )
else:
    self.proj_out = zero_module(nn.Linear(inner_dim, in_channels))

Temporal Transformer（VideoTransformer）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
self.time_stack = nn.ModuleList(
    [
        VideoTransformerBlock(
            inner_dim,
            n_time_mix_heads,
            time_mix_d_head,
            dropout=dropout,
            context_dim=time_context_dim,  # 1024
            timesteps=timesteps,
            checkpoint=checkpoint,
            ff_in=ff_in,
            inner_dim=time_mix_inner_dim,
            attn_mode=attn_mode,
            disable_self_attn=disable_self_attn,  # False
            disable_temporal_crossattention=disable_temporal_crossattention,  # False
        )
        for _ in range(self.depth)
    ]
)

Downsample3D

如果 resblock_updown 为 True，则添加 VideoResBlock / TemporalSpatialResBlock：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
self.input_blocks.append(
    TimestepEmbedSequential(
        get_resblock(
            merge_factor=merge_factor,
            merge_strategy=merge_strategy,
            video_kernel_size=video_kernel_size,
            ch=ch,
            time_embed_dim=time_embed_dim,
            dropout=dropout,
            out_ch=out_ch,
            dims=dims,
            use_checkpoint=use_checkpoint,
            use_scale_shift_norm=use_scale_shift_norm,
            down=True,  # downsample
        )
        if resblock_updown  # False
        else Downsample(  # downsample block
            ch,
            conv_resample,
            dims=dims,
            out_channels=out_ch,
            third_down=time_downup,
        )
    )
)

DownsampleBlock 是一个 kernel size = 3 的 2D 卷积网络（图中标注的维度可能不正确）：

MiddleResBlock

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
self.middle_block = TimestepEmbedSequential(
    get_resblock(
        merge_factor=merge_factor,
        merge_strategy=merge_strategy,
        video_kernel_size=video_kernel_size,
        ch=ch,
        time_embed_dim=time_embed_dim,
        out_ch=None,
        dropout=dropout,
        dims=dims,
        use_checkpoint=use_checkpoint,
        use_scale_shift_norm=use_scale_shift_norm,
    ),
    get_attention_layer(
        ch,
        num_heads,
        dim_head,
        depth=transformer_depth_middle,
        context_dim=context_dim,
        use_checkpoint=use_checkpoint,
    ),
    get_resblock(
        merge_factor=merge_factor,
        merge_strategy=merge_strategy,
        video_kernel_size=video_kernel_size,
        ch=ch,
        out_ch=None,
        time_embed_dim=time_embed_dim,
        dropout=dropout,
        dims=dims,
        use_checkpoint=use_checkpoint,
        use_scale_shift_norm=use_scale_shift_norm,
    ),
)

OutBlock

1
2
3
4
5
self.out = nn.Sequential(
    normalization(ch),
    nn.SiLU(),
    zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
)

Conditioner

CLIP

  • ViT 处理图像
  • Token embedding 处理文本输入

Overall Network Architecture

Training Details

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def fit(
    self,
    model: "pl.LightningModule",
    train_dataloaders: Optional[Union[TRAIN_DATALOADERS, LightningDataModule]] = None,
    val_dataloaders: Optional[EVAL_DATALOADERS] = None,
    datamodule: Optional[LightningDataModule] = None,
    train_dataloader=None,  # TODO: remove with 1.6
    ckpt_path: Optional[str] = None,
) -> None:
    r"""
    Runs the full optimization routine.

    Args:
        model: Model to fit.

        train_dataloaders: A collection of :class:`torch.utils.data.DataLoader` or a
            :class:`~pytorch_lightning.core.datamodule.LightningDataModule` specifying
            training samples. In the case of multiple dataloaders, please see this
            :ref:`page <multiple-training-dataloaders>`.

        val_dataloaders: A :class:`torch.utils.data.DataLoader` or a sequence of them
            specifying validation samples.

        ckpt_path: Path/URL of the checkpoint from which training is resumed. If there is
            no checkpoint file at the path, an exception is raised. If resuming from
            mid-epoch checkpoint, training will start from the beginning of the next epoch.

        datamodule: An instance of
            :class:`~pytorch_lightning.core.datamodule.LightningDataModule`.
    """
    if train_dataloader is not None:
        rank_zero_deprecation(
            "`trainer.fit(train_dataloader)` is deprecated in v1.4 and will be removed in v1.6."
            " Use `trainer.fit(train_dataloaders)` instead. HINT: added 's'"
        )
        train_dataloaders = train_dataloader
    self._call_and_handle_interrupt(
        self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
    )

Step 1：读取数据

1
2
3
4
5
6
7
8
9
10
11
12
13
def train_dataloader(self):  # training dataset dataloader
    sampler = DistributedSampler(self.train_dataset, seed=self.seed)
    return wds.WebLoader(self.train_dataset, batch_size=self.batch_size,
                         num_workers=self.num_workers, shuffle=False, sampler=sampler)

def val_dataloader(self):
    loader = wds.WebLoader(self.val_dataset, batch_size=self.batch_size,
                           num_workers=self.num_workers, shuffle=False)
    return loader

def test_dataloader(self):  # test dataset is the same as the evaluation dataset
    return wds.WebLoader(self.val_dataset, batch_size=self.batch_size,
                         num_workers=self.num_workers, shuffle=False)

Step 2：配置 optimizer

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
def configure_optimizers(self):
    lr = self.learning_rate
    if 'all' in self.trained_param_keys:
        params = list(self.model.parameters())  # all model parameters
    else:
        names = []
        params = []
        for name, param in self.model.named_parameters():
            flag = False
            for k in self.trained_param_keys:
                if k in name:
                    names += [name]
                    params += [param]
                    flag = True
                if flag:
                    break
        print(names)

    for embedder in self.conditioner.embedders:  # all conditioners are not trainable
        if embedder.is_trainable:
            params = params + list(embedder.parameters())

    opt = self.instantiate_optimizer_from_config(params, lr, self.optimizer_config)
    if self.scheduler_config is not None:
        scheduler = instantiate_from_config(self.scheduler_config)
        print("Setting up LambdaLR scheduler...")
        scheduler = [
            {
                "scheduler": LambdaLR(opt, lr_lambda=scheduler.schedule),
                "interval": "step",
                "frequency": 1,
            }
        ]
        return [opt], scheduler
    return opt

Step 3：两个关键 condition

  • cond_aug：加到每一帧上的噪声量，用于 cascade LDM 生成高分辨率结果。
  • motion bucket ID：数值越大，生成视频中的运动越多。

First stage encoding（VAE encode）

1
x = self.encode_first_stage(frames_reshape)  # get latent representation, (b t) c h w

Shape：\((B\times N) \times C_{vae} \times H \times W\)

添加 custom condition

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
def add_custom_cond(self, batch, infer=False):
    batch['num_video_frames'] = self.num_samples  # 16

    image = batch['video'][:, :, 0]  # 取第一帧图像作为 condition
    batch['cond_frames_without_noise'] = image.half()  # 转成 float16

    N = batch['video'].shape[0]  # B x C x N_frames x H x W，这里 N 即 batch size
    if not infer:  # 训练时，cond_aug 是随机的
        cond_aug = ((-3.0) + (0.5) * torch.randn((N,))).exp().cuda().half()
    else:          # 推理时，cond_aug 固定
        cond_aug = torch.full((N, ), 0.02).cuda().half()
    batch['cond_aug'] = cond_aug
    batch['cond_frames'] = (
        image + rearrange(cond_aug, 'b -> b 1 1 1') * torch.randn_like(image)
    ).half()

    # for dataset without indicator
    if not 'image_only_indicator' in batch:
        batch['image_only_indicator'] = torch.zeros((N, self.num_samples)).cuda().half()
    return batch

CLIP 编码第一帧作为 condition

1
2
3
4
5
6
def forward(self, vid):
    # 这里 vid 代表 batch['cond_frames_without_noise']，即第一帧图像
    vid = self.open_clip(vid)
    vid = rearrange(vid, "(b t) d -> b t d", t=self.n_cond_frames)
    vid = repeat(vid, "b t d -> (b s) t d", s=self.n_copies)
    return vid

Shape：\(1\times 1024 \rightarrow 1 \times 1 \times 1024 \rightarrow 1 \times 1 \times 1024\)

Aesthetic score

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def forward(self, x):
    B, C, T, H, W = x.shape

    y = x[:, :, T//2]
    y = torch.nn.functional.interpolate(y, [224, 384], mode='bilinear')
    y = y[:, :, :, 80:304]
    y = (y + 1) * 0.5
    y[:, 0] = (y[:, 0] - 0.48145466) / 0.26862954
    y[:, 1] = (y[:, 1] - 0.4578275) / 0.26130258
    y[:, 2] = (y[:, 2] - 0.40821073) / 0.27577711

    image_features = self.aesthetic_model.encode_image(y)             # CLIP image encoder, 1 x 768
    im_emb_arr = normalized(image_features.cpu().detach().numpy())    # 1 x 768
    aesthetic = self.aesthetic_mlp(
        torch.from_numpy(im_emb_arr).to('cuda').type(torch.cuda.FloatTensor)
    )                                                                # 1 x 1

    # [aesthetic_score, time_step]
    return torch.cat([aesthetic, timestep_embedding(aesthetic[:, 0] * 100, 255)], 1)

Shape：\(1 \times 256\)

Elevation

elevation（仰角）走的是 ConcatTimestepEmbedderND——它把输入的每一个维度独立做 timestep embedding，然后拼接起来：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
class ConcatTimestepEmbedderND(AbstractEmbModel):
    """embeds each dimension independently and concatenates them"""

    def __init__(self, outdim):
        super().__init__()
        self.timestep = Timestep(outdim)
        self.outdim = outdim

    def forward(self, x):
        if x.ndim == 1:
            x = x[:, None]                  # 标量 → (b, 1)
        assert len(x.shape) == 2
        b, dims = x.shape[0], x.shape[1]
        x = rearrange(x, "b d -> (b d)")    # 摊平，让所有标量一次过 Timestep
        emb = self.timestep(x)              # (b*d, outdim)，正弦位置编码
        emb = rearrange(emb, "(b d) d2 -> b (d d2)",
                        b=b, d=dims, d2=self.outdim)
        return emb

逻辑是：把 2D 输入的每一列都当作一个独立标量去编码。先把 batch 和维度一起摊平，让 Timestep 一次处理完所有标量，再拆回来，把各维度的 embedding（每个长度 outdim）沿特征轴拼接，最终输出形状为 \(b \times (\text{dims} \cdot \text{outdim})\)。

对 elevation 而言 dims = 1，所以输出就是 \(b \times \text{outdim}\)。同一个类也被 cond_aug、fps_id、motion_bucket_id 等标量 condition 复用——这就是为什么 SVD 里所有标量条件都能走同一套 embedder。

代码取自 Stability-AI generative-models 的 sgm/modules/encoders/modules.py。