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