| 503 | |
| 504 | |
| 505 | class Encoder3d(nn.Module): |
| 506 | |
| 507 | def __init__( |
| 508 | self, |
| 509 | dim=128, |
| 510 | z_dim=4, |
| 511 | dim_mult=[1, 2, 4, 4], |
| 512 | num_res_blocks=2, |
| 513 | attn_scales=[], |
| 514 | temperal_downsample=[True, True, False], |
| 515 | dropout=0.0, |
| 516 | ): |
| 517 | super().__init__() |
| 518 | self.dim = dim |
| 519 | self.z_dim = z_dim |
| 520 | self.dim_mult = dim_mult |
| 521 | self.num_res_blocks = num_res_blocks |
| 522 | self.attn_scales = attn_scales |
| 523 | self.temperal_downsample = temperal_downsample |
| 524 | |
| 525 | # dimensions |
| 526 | dims = [dim * u for u in [1] + dim_mult] |
| 527 | scale = 1.0 |
| 528 | |
| 529 | # init block |
| 530 | self.conv1 = CausalConv3d(12, dims[0], 3, padding=1) |
| 531 | |
| 532 | # downsample blocks |
| 533 | downsamples = [] |
| 534 | for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): |
| 535 | t_down_flag = ( |
| 536 | temperal_downsample[i] |
| 537 | if i < len(temperal_downsample) else False) |
| 538 | downsamples.append( |
| 539 | Down_ResidualBlock( |
| 540 | in_dim=in_dim, |
| 541 | out_dim=out_dim, |
| 542 | dropout=dropout, |
| 543 | mult=num_res_blocks, |
| 544 | temperal_downsample=t_down_flag, |
| 545 | down_flag=i != len(dim_mult) - 1, |
| 546 | )) |
| 547 | scale /= 2.0 |
| 548 | self.downsamples = nn.Sequential(*downsamples) |
| 549 | |
| 550 | # middle blocks |
| 551 | self.middle = nn.Sequential( |
| 552 | ResidualBlock(out_dim, out_dim, dropout), |
| 553 | AttentionBlock(out_dim), |
| 554 | ResidualBlock(out_dim, out_dim, dropout), |
| 555 | ) |
| 556 | |
| 557 | # # output blocks |
| 558 | self.head = nn.Sequential( |
| 559 | RMS_norm(out_dim, images=False), |
| 560 | nn.SiLU(), |
| 561 | CausalConv3d(out_dim, z_dim, 3, padding=1), |
| 562 | ) |