wanvae ¶

def __init__(
    self,
    config: WanVAEConfig,
) -> None:
    nn.Module.__init__(self)
    ParallelTiledVAE.__init__(self, config)

    self.z_dim = config.z_dim
    self.temperal_downsample = list(config.temperal_downsample)
    self.temperal_upsample = list(config.temperal_downsample)[::-1]

    decoder_base_dim = config.base_dim if config.decoder_base_dim is None else config.decoder_base_dim

    self.latents_mean = list(config.latents_mean)
    self.latents_std = list(config.latents_std)
    self.shift_factor = config.shift_factor

    if config.load_encoder:
        self.encoder = WanEncoder3d(
            in_channels=config.in_channels,
            dim=config.base_dim, 
            z_dim=self.z_dim * 2,
            dim_mult=config.dim_mult, 
            num_res_blocks=config.num_res_blocks,
            attn_scales=config.attn_scales,
            temperal_downsample=self.temperal_downsample,
            dropout=config.dropout,
            is_residual=config.is_residual,
        )
    self.quant_conv = WanCausalConv3d(self.z_dim * 2, self.z_dim * 2, 1)
    self.post_quant_conv = WanCausalConv3d(self.z_dim, self.z_dim, 1)

    if config.load_decoder:
        self.decoder = WanDecoder3d(
            dim=decoder_base_dim, 
            z_dim=self.z_dim,
            dim_mult=config.dim_mult, 
            num_res_blocks=config.num_res_blocks,
            attn_scales=config.attn_scales,
            temperal_upsample=self.temperal_upsample, 
            dropout=config.dropout,
            out_channels=config.out_channels,
            is_residual=config.is_residual,
        )

    self.use_feature_cache = config.use_feature_cache

def __init__(self, dim) -> None:
    super().__init__()
    self.dim = dim

    # layers
    self.norm = WanRMS_norm(dim)
    self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
    self.proj = nn.Conv2d(dim, dim, 1)

def __init__(
    self,
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, int, int],
    stride: int | tuple[int, int, int] = 1,
    padding: int | tuple[int, int, int] = 0,
) -> None:
    super().__init__(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=kernel_size,
        stride=stride,
        padding=padding,
    )
    self.padding: tuple[int, int, int]
    # Set up causal padding
    self._padding: tuple[int, ...] = (self.padding[2], self.padding[2],
                                      self.padding[1], self.padding[1],
                                      2 * self.padding[0], 0)
    self.padding = (0, 0, 0)

def __init__(
    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,
):
    super().__init__()
    self.dim = dim
    self.z_dim = z_dim
    dim_mult = list(dim_mult)
    self.dim_mult = dim_mult
    self.num_res_blocks = num_res_blocks
    self.attn_scales = list(attn_scales)
    self.temperal_upsample = list(temperal_upsample)

    self.nonlinearity = get_act_fn(non_linearity)

    # dimensions
    dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]

    # init block
    self.conv_in = WanCausalConv3d(z_dim, dims[0], 3, padding=1)

    # middle blocks
    self.mid_block = WanMidBlock(dims[0],
                                 dropout,
                                 non_linearity,
                                 num_layers=1)

    # upsample blocks
    self.up_blocks = nn.ModuleList([])
    for i, (in_dim,
            out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)):
        # residual (+attention) blocks
        if i > 0 and not is_residual:
            # wan vae 2.1
            in_dim = in_dim // 2

        # determine if we need upsampling
        up_flag = i != len(dim_mult) - 1
        # determine upsampling mode, if not upsampling, set to None
        upsample_mode = None
        if up_flag and temperal_upsample[i]:
            upsample_mode = "upsample3d"
        elif up_flag:
            upsample_mode = "upsample2d"

        # Create and add the upsampling block
        if is_residual:
            up_block = WanResidualUpBlock(
                in_dim=in_dim,
                out_dim=out_dim,
                num_res_blocks=num_res_blocks,
                dropout=dropout,
                temperal_upsample=temperal_upsample[i] if up_flag else False,
                up_flag=up_flag,
                non_linearity=non_linearity,
            )
        else:
            up_block = WanUpBlock(
                in_dim=in_dim,
                out_dim=out_dim,
                num_res_blocks=num_res_blocks,
                dropout=dropout,
                upsample_mode=upsample_mode,
                non_linearity=non_linearity,
            )
        self.up_blocks.append(up_block)

    # output blocks
    self.norm_out = WanRMS_norm(out_dim, images=False)
    self.conv_out = WanCausalConv3d(out_dim, out_channels, 3, padding=1)

    self.gradient_checkpointing = False

def __init__(
    self,
    in_channels: int = 3,
    dim=128,
    z_dim=4,
    dim_mult=(1, 2, 4, 4),
    num_res_blocks=2,
    attn_scales=(),
    temperal_downsample=(True, True, False),
    dropout=0.0,
    non_linearity: str = "silu",
    is_residual: bool = False,  # wan 2.2 vae use a residual downblock
):
    super().__init__()
    self.dim = dim
    self.z_dim = z_dim
    dim_mult = list(dim_mult)
    self.dim_mult = dim_mult
    self.num_res_blocks = num_res_blocks
    self.attn_scales = list(attn_scales)
    self.temperal_downsample = list(temperal_downsample)
    self.nonlinearity = get_act_fn(non_linearity)

    # dimensions
    dims = [dim * u for u in [1] + dim_mult]
    scale = 1.0

    # init block
    self.conv_in = WanCausalConv3d(in_channels, dims[0], 3, padding=1)

    # downsample blocks
    self.down_blocks = nn.ModuleList([])
    for i, (in_dim,
            out_dim) in enumerate(zip(dims[:-1], dims[1:], strict=True)):
        # residual (+attention) blocks
        if is_residual:
            self.down_blocks.append(
                WanResidualDownBlock(
                    in_dim,
                    out_dim,
                    dropout,
                    num_res_blocks,
                    temperal_downsample=temperal_downsample[i] if i != len(dim_mult) - 1 else False,
                    down_flag=i != len(dim_mult) - 1,
                )
            )
        else:
            for _ in range(num_res_blocks):
                self.down_blocks.append(WanResidualBlock(in_dim, out_dim, dropout))
                if scale in attn_scales:
                    self.down_blocks.append(WanAttentionBlock(out_dim))
                in_dim = out_dim

            # downsample block
            if i != len(dim_mult) - 1:
                mode = "downsample3d" if temperal_downsample[i] else "downsample2d"
                self.down_blocks.append(WanResample(out_dim, mode=mode))
                scale /= 2.0

    # middle blocks
    self.mid_block = WanMidBlock(out_dim,
                                 dropout,
                                 non_linearity,
                                 num_layers=1)

    # output blocks
    self.norm_out = WanRMS_norm(out_dim, images=False)
    self.conv_out = WanCausalConv3d(out_dim, z_dim, 3, padding=1)

    self.gradient_checkpointing = False

def __init__(self,
             dim: int,
             dropout: float = 0.0,
             non_linearity: str = "silu",
             num_layers: int = 1):
    super().__init__()
    self.dim = dim

    # Create the components
    resnets = [WanResidualBlock(dim, dim, dropout, non_linearity)]
    attentions = []
    for _ in range(num_layers):
        attentions.append(WanAttentionBlock(dim))
        resnets.append(WanResidualBlock(dim, dim, dropout, non_linearity))
    self.attentions = nn.ModuleList(attentions)
    self.resnets = nn.ModuleList(resnets)

    self.gradient_checkpointing = False

def __init__(self,
             dim: int,
             channel_first: bool = True,
             images: bool = True,
             bias: bool = False) -> None:
    super().__init__()
    broadcastable_dims = (1, 1, 1) if not images else (1, 1)
    shape = (dim, *broadcastable_dims) if channel_first else (dim, )

    self.channel_first = channel_first
    self.scale = dim**0.5
    self.gamma = nn.Parameter(torch.ones(shape))
    self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0

def __init__(self, dim: int, mode: str, upsample_out_dim: int = None) -> None:
    super().__init__()
    self.dim = dim
    self.mode = mode

    # default to dim //2
    if upsample_out_dim is None:
        upsample_out_dim = dim // 2

    # layers
    if mode == "upsample2d":
        self.resample = nn.Sequential(
            WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
            nn.Conv2d(dim, upsample_out_dim, 3, padding=1),
        )
    elif mode == "upsample3d":
        self.resample = nn.Sequential(
            WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
            nn.Conv2d(dim, upsample_out_dim, 3, padding=1),
        )
        self.time_conv = WanCausalConv3d(dim,
                                         dim * 2, (3, 1, 1),
                                         padding=(1, 0, 0))

    elif mode == "downsample2d":
        self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)),
                                      nn.Conv2d(dim, dim, 3, stride=(2, 2)))
    elif mode == "downsample3d":
        self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)),
                                      nn.Conv2d(dim, dim, 3, stride=(2, 2)))
        self.time_conv = WanCausalConv3d(dim,
                                         dim, (3, 1, 1),
                                         stride=(2, 1, 1),
                                         padding=(0, 0, 0))

    else:
        self.resample = nn.Identity()

def __init__(
    self,
    in_dim: int,
    out_dim: int,
    dropout: float = 0.0,
    non_linearity: str = "silu",
) -> None:
    super().__init__()
    self.in_dim = in_dim
    self.out_dim = out_dim
    self.nonlinearity = get_act_fn(non_linearity)

    # layers
    self.norm1 = WanRMS_norm(in_dim, images=False)
    self.conv1 = WanCausalConv3d(in_dim, out_dim, 3, padding=1)
    self.norm2 = WanRMS_norm(out_dim, images=False)
    self.dropout = nn.Dropout(dropout)
    self.conv2 = WanCausalConv3d(out_dim, out_dim, 3, padding=1)
    self.conv_shortcut = WanCausalConv3d(
        in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()

def __init__(
    self,
    in_dim: int,
    out_dim: int,
    num_res_blocks: int,
    dropout: float = 0.0,
    temperal_upsample: bool = False,
    up_flag: bool = False,
    non_linearity: str = "silu",
):
    super().__init__()
    self.in_dim = in_dim
    self.out_dim = out_dim

    if up_flag:
        self.avg_shortcut = DupUp3D(
            in_dim,
            out_dim,
            factor_t=2 if temperal_upsample else 1,
            factor_s=2,
        )
    else:
        self.avg_shortcut = None

    # create residual blocks
    resnets = []
    current_dim = in_dim
    for _ in range(num_res_blocks + 1):
        resnets.append(WanResidualBlock(current_dim, out_dim, dropout, non_linearity))
        current_dim = out_dim

    self.resnets = nn.ModuleList(resnets)

    # Add upsampling layer if needed
    if up_flag:
        upsample_mode = "upsample3d" if temperal_upsample else "upsample2d"
        self.upsampler = WanResample(out_dim, mode=upsample_mode, upsample_out_dim=out_dim)
    else:
        self.upsampler = None

    self.gradient_checkpointing = False

def __init__(
    self,
    in_dim: int,
    out_dim: int,
    num_res_blocks: int,
    dropout: float = 0.0,
    upsample_mode: str | None = None,
    non_linearity: str = "silu",
):
    super().__init__()
    self.in_dim = in_dim
    self.out_dim = out_dim

    # Create layers list
    resnets = []
    # Add residual blocks and attention if needed
    current_dim = in_dim
    for _ in range(num_res_blocks + 1):
        resnets.append(
            WanResidualBlock(current_dim, out_dim, dropout, non_linearity))
        current_dim = out_dim

    self.resnets = nn.ModuleList(resnets)

    # Add upsampling layer if needed
    self.upsamplers = None
    if upsample_mode is not None:
        self.upsamplers = nn.ModuleList(
            [WanResample(out_dim, mode=upsample_mode)])

    self.gradient_checkpointing = False