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 upsample_mode (str, optional): Mode for upsampli
| 712 | |
| 713 | |
| 714 | class WanUpBlock(nn.Module): |
| 715 | """ |
| 716 | A block that handles upsampling for the WanVAE decoder. |
| 717 | |
| 718 | Args: |
| 719 | in_dim (int): Input dimension |
| 720 | out_dim (int): Output dimension |
| 721 | num_res_blocks (int): Number of residual blocks |
| 722 | dropout (float): Dropout rate |
| 723 | upsample_mode (str, optional): Mode for upsampling ('upsample2d' or 'upsample3d') |
| 724 | non_linearity (str): Type of non-linearity to use |
| 725 | """ |
| 726 | |
| 727 | def __init__( |
| 728 | self, |
| 729 | in_dim: int, |
| 730 | out_dim: int, |
| 731 | num_res_blocks: int, |
| 732 | dropout: float = 0.0, |
| 733 | upsample_mode: Optional[str] = None, |
| 734 | non_linearity: str = "silu", |
| 735 | ): |
| 736 | super().__init__() |
| 737 | self.in_dim = in_dim |
| 738 | self.out_dim = out_dim |
| 739 | |
| 740 | # Create layers list |
| 741 | resnets = [] |
| 742 | # Add residual blocks and attention if needed |
| 743 | current_dim = in_dim |
| 744 | for _ in range(num_res_blocks + 1): |
| 745 | resnets.append(WanResidualBlock(current_dim, out_dim, dropout, non_linearity)) |
| 746 | current_dim = out_dim |
| 747 | |
| 748 | self.resnets = nn.ModuleList(resnets) |
| 749 | |
| 750 | # Add upsampling layer if needed |
| 751 | self.upsamplers = None |
| 752 | if upsample_mode is not None: |
| 753 | self.upsamplers = nn.ModuleList([WanResample(out_dim, mode=upsample_mode)]) |
| 754 | |
| 755 | self.gradient_checkpointing = False |
| 756 | |
| 757 | def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=None): |
| 758 | """ |
| 759 | Forward pass through the upsampling block. |
| 760 | |
| 761 | Args: |
| 762 | x (torch.Tensor): Input tensor |
| 763 | feat_cache (list, optional): Feature cache for causal convolutions |
| 764 | feat_idx (list, optional): Feature index for cache management |
| 765 | |
| 766 | Returns: |
| 767 | torch.Tensor: Output tensor |
| 768 | """ |
| 769 | for resnet in self.resnets: |
| 770 | if feat_cache is not None: |
| 771 | x = resnet(x, feat_cache, feat_idx) |