| 515 | |
| 516 | |
| 517 | class Encoder3d(nn.Module): |
| 518 | |
| 519 | def __init__(self, |
| 520 | dim=128, |
| 521 | z_dim=4, |
| 522 | dim_mult=[1, 2, 4, 4], |
| 523 | num_res_blocks=2, |
| 524 | attn_scales=[], |
| 525 | temperal_downsample=[True, True, False], |
| 526 | dropout=0.0): |
| 527 | super().__init__() |
| 528 | self.dim = dim |
| 529 | self.z_dim = z_dim |
| 530 | self.dim_mult = dim_mult |
| 531 | self.num_res_blocks = num_res_blocks |
| 532 | self.attn_scales = attn_scales |
| 533 | self.temperal_downsample = temperal_downsample |
| 534 | |
| 535 | # dimensions |
| 536 | dims = [dim * u for u in [1] + dim_mult] |
| 537 | scale = 1.0 |
| 538 | |
| 539 | # init block |
| 540 | self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) |
| 541 | |
| 542 | # downsample blocks |
| 543 | downsamples = [] |
| 544 | for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): |
| 545 | # residual (+attention) blocks |
| 546 | for _ in range(num_res_blocks): |
| 547 | downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) |
| 548 | if scale in attn_scales: |
| 549 | downsamples.append(AttentionBlock(out_dim)) |
| 550 | in_dim = out_dim |
| 551 | |
| 552 | # downsample block |
| 553 | if i != len(dim_mult) - 1: |
| 554 | mode = 'downsample3d' if temperal_downsample[ |
| 555 | i] else 'downsample2d' |
| 556 | downsamples.append(Resample(out_dim, mode=mode)) |
| 557 | scale /= 2.0 |
| 558 | self.downsamples = nn.Sequential(*downsamples) |
| 559 | |
| 560 | # middle blocks |
| 561 | self.middle = nn.Sequential(ResidualBlock(out_dim, out_dim, dropout), |
| 562 | AttentionBlock(out_dim), |
| 563 | ResidualBlock(out_dim, out_dim, dropout)) |
| 564 | |
| 565 | # output blocks |
| 566 | self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), |
| 567 | CausalConv3d(out_dim, z_dim, 3, padding=1)) |
| 568 | |
| 569 | def forward(self, x, feat_cache=None, feat_idx=[0]): |
| 570 | if feat_cache is not None: |
| 571 | idx = feat_idx[0] |
| 572 | cache_x = x[:, :, -CACHE_T:, :, :].clone() |
| 573 | if cache_x.shape[2] < 2 and feat_cache[idx] is not None: |
| 574 | # cache last frame of last two chunk |