| 467 | |
| 468 | |
| 469 | class WanResidualDownBlock(nn.Module): |
| 470 | def __init__(self, in_dim, out_dim, dropout, num_res_blocks, temperal_downsample=False, down_flag=False): |
| 471 | super().__init__() |
| 472 | |
| 473 | # Shortcut path with downsample |
| 474 | self.avg_shortcut = AvgDown3D( |
| 475 | in_dim, |
| 476 | out_dim, |
| 477 | factor_t=2 if temperal_downsample else 1, |
| 478 | factor_s=2 if down_flag else 1, |
| 479 | ) |
| 480 | |
| 481 | # Main path with residual blocks and downsample |
| 482 | resnets = [] |
| 483 | for _ in range(num_res_blocks): |
| 484 | resnets.append(WanResidualBlock(in_dim, out_dim, dropout)) |
| 485 | in_dim = out_dim |
| 486 | self.resnets = nn.ModuleList(resnets) |
| 487 | |
| 488 | # Add the final downsample block |
| 489 | if down_flag: |
| 490 | mode = "downsample3d" if temperal_downsample else "downsample2d" |
| 491 | self.downsampler = WanResample(out_dim, mode=mode) |
| 492 | else: |
| 493 | self.downsampler = None |
| 494 | |
| 495 | def forward(self, x, feat_cache=None, feat_idx=[0]): |
| 496 | x_copy = x.clone() |
| 497 | for resnet in self.resnets: |
| 498 | x = resnet(x, feat_cache, feat_idx) |
| 499 | if self.downsampler is not None: |
| 500 | x = self.downsampler(x, feat_cache, feat_idx) |
| 501 | |
| 502 | return x + self.avg_shortcut(x_copy) |
| 503 | |
| 504 | |
| 505 | class WanEncoder3d(nn.Module): |