(
self,
dim=128,
z_dim=4,
dim_mult=[1, 2, 4, 4],
num_res_blocks=2,
attn_scales=[],
temperal_upsample=[False, True, True],
dropout=0.0,
non_linearity: str = "silu",
out_channels: int = 3,
is_residual: bool = False,
)
| 796 | """ |
| 797 | |
| 798 | def __init__( |
| 799 | self, |
| 800 | dim=128, |
| 801 | z_dim=4, |
| 802 | dim_mult=[1, 2, 4, 4], |
| 803 | num_res_blocks=2, |
| 804 | attn_scales=[], |
| 805 | temperal_upsample=[False, True, True], |
| 806 | dropout=0.0, |
| 807 | non_linearity: str = "silu", |
| 808 | out_channels: int = 3, |
| 809 | is_residual: bool = False, |
| 810 | ): |
| 811 | super().__init__() |
| 812 | self.dim = dim |
| 813 | self.z_dim = z_dim |
| 814 | self.dim_mult = dim_mult |
| 815 | self.num_res_blocks = num_res_blocks |
| 816 | self.attn_scales = attn_scales |
| 817 | self.temperal_upsample = temperal_upsample |
| 818 | |
| 819 | self.nonlinearity = get_activation(non_linearity) |
| 820 | |
| 821 | # dimensions |
| 822 | dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] |
| 823 | |
| 824 | # init block |
| 825 | self.conv_in = WanCausalConv3d(z_dim, dims[0], 3, padding=1) |
| 826 | |
| 827 | # middle blocks |
| 828 | self.mid_block = WanMidBlock(dims[0], dropout, non_linearity, num_layers=1) |
| 829 | |
| 830 | # upsample blocks |
| 831 | self.up_blocks = nn.ModuleList([]) |
| 832 | for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): |
| 833 | # residual (+attention) blocks |
| 834 | if i > 0 and not is_residual: |
| 835 | # wan vae 2.1 |
| 836 | in_dim = in_dim // 2 |
| 837 | |
| 838 | # determine if we need upsampling |
| 839 | up_flag = i != len(dim_mult) - 1 |
| 840 | # determine upsampling mode, if not upsampling, set to None |
| 841 | upsample_mode = None |
| 842 | if up_flag and temperal_upsample[i]: |
| 843 | upsample_mode = "upsample3d" |
| 844 | elif up_flag: |
| 845 | upsample_mode = "upsample2d" |
| 846 | # Create and add the upsampling block |
| 847 | if is_residual: |
| 848 | up_block = WanResidualUpBlock( |
| 849 | in_dim=in_dim, |
| 850 | out_dim=out_dim, |
| 851 | num_res_blocks=num_res_blocks, |
| 852 | dropout=dropout, |
| 853 | temperal_upsample=temperal_upsample[i] if up_flag else False, |
| 854 | up_flag=up_flag, |
| 855 | non_linearity=non_linearity, |
nothing calls this directly
no test coverage detected