| 371 | |
| 372 | |
| 373 | class Decoder3d(nn.Module): |
| 374 | |
| 375 | def __init__(self, |
| 376 | dim=128, |
| 377 | z_dim=4, |
| 378 | dim_mult=[1, 2, 4, 4], |
| 379 | num_res_blocks=2, |
| 380 | attn_scales=[], |
| 381 | temperal_upsample=[False, True, True], |
| 382 | dropout=0.0): |
| 383 | super().__init__() |
| 384 | self.dim = dim |
| 385 | self.z_dim = z_dim |
| 386 | self.dim_mult = dim_mult |
| 387 | self.num_res_blocks = num_res_blocks |
| 388 | self.attn_scales = attn_scales |
| 389 | self.temperal_upsample = temperal_upsample |
| 390 | |
| 391 | # dimensions |
| 392 | dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] |
| 393 | scale = 1.0 / 2**(len(dim_mult) - 2) |
| 394 | |
| 395 | # init block |
| 396 | self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) |
| 397 | |
| 398 | # middle blocks |
| 399 | self.middle = nn.Sequential( |
| 400 | ResidualBlock(dims[0], dims[0], dropout), AttentionBlock(dims[0]), |
| 401 | ResidualBlock(dims[0], dims[0], dropout)) |
| 402 | |
| 403 | # upsample blocks |
| 404 | upsamples = [] |
| 405 | for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): |
| 406 | # residual (+attention) blocks |
| 407 | if i == 1 or i == 2 or i == 3: |
| 408 | in_dim = in_dim // 2 |
| 409 | for _ in range(num_res_blocks + 1): |
| 410 | upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) |
| 411 | if scale in attn_scales: |
| 412 | upsamples.append(AttentionBlock(out_dim)) |
| 413 | in_dim = out_dim |
| 414 | |
| 415 | # upsample block |
| 416 | if i != len(dim_mult) - 1: |
| 417 | mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d' |
| 418 | upsamples.append(Resample(out_dim, mode=mode)) |
| 419 | scale *= 2.0 |
| 420 | self.upsamples = nn.Sequential(*upsamples) |
| 421 | |
| 422 | # output blocks |
| 423 | self.head = nn.Sequential( |
| 424 | RMS_norm(out_dim, images=False), nn.SiLU(), |
| 425 | CausalConv3d(out_dim, 3, 3, padding=1)) |
| 426 | |
| 427 | def forward(self, x, feat_cache=None, feat_idx=[0]): |
| 428 | ## conv1 |
| 429 | if feat_cache is not None: |
| 430 | idx = feat_idx[0] |