A block that handles upsampling for the WanVAE decoder. Args: in_dim (int): Input dimension out_dim (int): Output dimension num_res_blocks (int): Number of residual blocks dropout (float): Dropout rate temperal_upsample (bool): Whether to upsample on
| 624 | |
| 625 | |
| 626 | class WanResidualUpBlock(nn.Module): |
| 627 | """ |
| 628 | A block that handles upsampling for the WanVAE decoder. |
| 629 | |
| 630 | Args: |
| 631 | in_dim (int): Input dimension |
| 632 | out_dim (int): Output dimension |
| 633 | num_res_blocks (int): Number of residual blocks |
| 634 | dropout (float): Dropout rate |
| 635 | temperal_upsample (bool): Whether to upsample on temporal dimension |
| 636 | up_flag (bool): Whether to upsample or not |
| 637 | non_linearity (str): Type of non-linearity to use |
| 638 | """ |
| 639 | |
| 640 | def __init__( |
| 641 | self, |
| 642 | in_dim: int, |
| 643 | out_dim: int, |
| 644 | num_res_blocks: int, |
| 645 | dropout: float = 0.0, |
| 646 | temperal_upsample: bool = False, |
| 647 | up_flag: bool = False, |
| 648 | non_linearity: str = "silu", |
| 649 | ): |
| 650 | super().__init__() |
| 651 | self.in_dim = in_dim |
| 652 | self.out_dim = out_dim |
| 653 | |
| 654 | if up_flag: |
| 655 | self.avg_shortcut = DupUp3D( |
| 656 | in_dim, |
| 657 | out_dim, |
| 658 | factor_t=2 if temperal_upsample else 1, |
| 659 | factor_s=2, |
| 660 | ) |
| 661 | else: |
| 662 | self.avg_shortcut = None |
| 663 | |
| 664 | # create residual blocks |
| 665 | resnets = [] |
| 666 | current_dim = in_dim |
| 667 | for _ in range(num_res_blocks + 1): |
| 668 | resnets.append(WanResidualBlock(current_dim, out_dim, dropout, non_linearity)) |
| 669 | current_dim = out_dim |
| 670 | |
| 671 | self.resnets = nn.ModuleList(resnets) |
| 672 | |
| 673 | # Add upsampling layer if needed |
| 674 | if up_flag: |
| 675 | upsample_mode = "upsample3d" if temperal_upsample else "upsample2d" |
| 676 | self.upsampler = WanResample(out_dim, mode=upsample_mode, upsample_out_dim=out_dim) |
| 677 | else: |
| 678 | self.upsampler = None |
| 679 | |
| 680 | self.gradient_checkpointing = False |
| 681 | |
| 682 | def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): |
| 683 | """ |