Update custom_pipeline/AutoencoderKLWan.py

custom_pipeline/AutoencoderKLWan.py

@@ -1,17 +1,3 @@
-# Copyright 2025 The Wan Team and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 from typing import List, Optional, Tuple, Union
 
 import torch
@@ -19,28 +5,24 @@ import torch.nn as nn
 import torch.nn.functional as F
 import torch.utils.checkpoint
 
-from ...configuration_utils import ConfigMixin, register_to_config
-from ...loaders import FromOriginalModelMixin
-from ...utils import logging
-from ...utils.accelerate_utils import apply_forward_hook
-from ..activations import get_activation
-from ..modeling_outputs import AutoencoderKLOutput
-from ..modeling_utils import ModelMixin
-from .vae import DecoderOutput, DiagonalGaussianDistribution
-
+from diffusers.configuration_utils import ConfigMixin, register_to_config  # 상대 → 절대 임포트
+from diffusers.loaders import FromOriginalModelMixin  # 상대 → 절대 임포트
+from diffusers.utils import logging  # 상대 → 절대 임포트
+from diffusers.utils.accelerate_utils import apply_forward_hook  # 상대 → 절대 임포트
+from diffusers.models.activations import get_activation  # 상대 → 절대 임포트
+from diffusers.models.modeling_outputs import AutoencoderKLOutput  # 상대 → 절대 임포트
+from diffusers.models.modeling_utils import ModelMixin  # 상대 → 절대 임포트
+from diffusers.models.vae import DecoderOutput, DiagonalGaussianDistribution  # 상대 → 절대 임포트
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 CACHE_T = 2
 
-
 class WanCausalConv3d(nn.Conv3d):
     r"""
     A custom 3D causal convolution layer with feature caching support.
-
     This layer extends the standard Conv3D layer by ensuring causality in the time dimension and handling feature
     caching for efficient inference.
-
     Args:
         in_channels (int): Number of channels in the input image
         out_channels (int): Number of channels produced by the convolution
@@ -48,7 +30,6 @@ class WanCausalConv3d(nn.Conv3d):
         stride (int or tuple, optional): Stride of the convolution. Default: 1
         padding (int or tuple, optional): Zero-padding added to all three sides of the input. Default: 0
     """
-
     def __init__(
         self,
         in_channels: int,
@@ -64,8 +45,6 @@ class WanCausalConv3d(nn.Conv3d):
             stride=stride,
             padding=padding,
         )
-
-        # Set up causal padding
         self._padding = (self.padding[2], self.padding[2], self.padding[1], self.padding[1], 2 * self.padding[0], 0)
         self.padding = (0, 0, 0)
 
@@ -78,11 +57,9 @@ class WanCausalConv3d(nn.Conv3d):
         x = F.pad(x, padding)
         return super().forward(x)
 
-
 class WanRMS_norm(nn.Module):
     r"""
     A custom RMS normalization layer.
-
     Args:
         dim (int): The number of dimensions to normalize over.
         channel_first (bool, optional): Whether the input tensor has channels as the first dimension.
@@ -90,12 +67,10 @@ class WanRMS_norm(nn.Module):
         images (bool, optional): Whether the input represents image data. Default is True.
         bias (bool, optional): Whether to include a learnable bias term. Default is 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))
@@ -104,26 +79,20 @@ class WanRMS_norm(nn.Module):
     def forward(self, x):
         return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias
 
-
 class WanUpsample(nn.Upsample):
     r"""
     Perform upsampling while ensuring the output tensor has the same data type as the input.
-
     Args:
         x (torch.Tensor): Input tensor to be upsampled.
-
     Returns:
         torch.Tensor: Upsampled tensor with the same data type as the input.
     """
-
     def forward(self, x):
         return super().forward(x.float()).type_as(x)
 
-
 class WanResample(nn.Module):
     r"""
     A custom resampling module for 2D and 3D data.
-
     Args:
         dim (int): The number of input/output channels.
         mode (str): The resampling mode. Must be one of:
@@ -133,13 +102,10 @@ class WanResample(nn.Module):
             - 'downsample2d': 2D downsampling with zero-padding and convolution.
             - 'downsample3d': 3D downsampling with zero-padding, convolution, and causal 3D convolution.
     """
-
     def __init__(self, dim: int, mode: str) -> None:
         super().__init__()
         self.dim = dim
         self.mode = mode
-
-        # layers
         if mode == "upsample2d":
             self.resample = nn.Sequential(
                 WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, dim // 2, 3, padding=1)
@@ -149,13 +115,11 @@ class WanResample(nn.Module):
                 WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, dim // 2, 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()
 
@@ -170,7 +134,6 @@ class WanResample(nn.Module):
                 else:
                     cache_x = x[:, :, -CACHE_T:, :, :].clone()
                     if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep":
-                        # cache last frame of last two chunk
                         cache_x = torch.cat(
                             [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2
                         )
@@ -182,7 +145,6 @@ class WanResample(nn.Module):
                         x = self.time_conv(x, feat_cache[idx])
                     feat_cache[idx] = cache_x
                     feat_idx[0] += 1
-
                     x = x.reshape(b, 2, c, t, h, w)
                     x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
                     x = x.reshape(b, c, t * 2, h, w)
@@ -190,7 +152,6 @@ class WanResample(nn.Module):
         x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
         x = self.resample(x)
         x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4)
-
         if self.mode == "downsample3d":
             if feat_cache is not None:
                 idx = feat_idx[0]
@@ -204,18 +165,15 @@ class WanResample(nn.Module):
                     feat_idx[0] += 1
         return x
 
-
 class WanResidualBlock(nn.Module):
     r"""
     A custom residual block module.
-
     Args:
         in_dim (int): Number of input channels.
         out_dim (int): Number of output channels.
         dropout (float, optional): Dropout rate for the dropout layer. Default is 0.0.
         non_linearity (str, optional): Type of non-linearity to use. Default is "silu".
     """
-
     def __init__(
         self,
         in_dim: int,
@@ -227,8 +185,6 @@ class WanResidualBlock(nn.Module):
         self.in_dim = in_dim
         self.out_dim = out_dim
         self.nonlinearity = get_activation(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)
@@ -237,61 +193,43 @@ class WanResidualBlock(nn.Module):
         self.conv_shortcut = WanCausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()
 
     def forward(self, x, feat_cache=None, feat_idx=[0]):
-        # Apply shortcut connection
         h = self.conv_shortcut(x)
-
-        # First normalization and activation
         x = self.norm1(x)
         x = self.nonlinearity(x)
-
         if feat_cache is not None:
             idx = feat_idx[0]
             cache_x = x[:, :, -CACHE_T:, :, :].clone()
             if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                 cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
-
             x = self.conv1(x, feat_cache[idx])
             feat_cache[idx] = cache_x
             feat_idx[0] += 1
         else:
             x = self.conv1(x)
-
-        # Second normalization and activation
         x = self.norm2(x)
         x = self.nonlinearity(x)
-
-        # Dropout
         x = self.dropout(x)
-
         if feat_cache is not None:
             idx = feat_idx[0]
             cache_x = x[:, :, -CACHE_T:, :, :].clone()
             if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                 cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
-
             x = self.conv2(x, feat_cache[idx])
             feat_cache[idx] = cache_x
             feat_idx[0] += 1
         else:
             x = self.conv2(x)
-
-        # Add residual connection
         return x + h
 
-
 class WanAttentionBlock(nn.Module):
     r"""
     Causal self-attention with a single head.
-
     Args:
         dim (int): The number of channels in the input tensor.
     """
-
     def __init__(self, dim):
         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)
@@ -299,46 +237,30 @@ class WanAttentionBlock(nn.Module):
     def forward(self, x):
         identity = x
         batch_size, channels, time, height, width = x.size()
-
         x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * time, channels, height, width)
         x = self.norm(x)
-
-        # compute query, key, value
         qkv = self.to_qkv(x)
         qkv = qkv.reshape(batch_size * time, 1, channels * 3, -1)
         qkv = qkv.permute(0, 1, 3, 2).contiguous()
         q, k, v = qkv.chunk(3, dim=-1)
-
-        # apply attention
         x = F.scaled_dot_product_attention(q, k, v)
-
         x = x.squeeze(1).permute(0, 2, 1).reshape(batch_size * time, channels, height, width)
-
-        # output projection
         x = self.proj(x)
-
-        # Reshape back: [(b*t), c, h, w] -> [b, c, t, h, w]
         x = x.view(batch_size, time, channels, height, width)
         x = x.permute(0, 2, 1, 3, 4)
-
         return x + identity
 
-
 class WanMidBlock(nn.Module):
     """
     Middle block for WanVAE encoder and decoder.
-
     Args:
         dim (int): Number of input/output channels.
         dropout (float): Dropout rate.
         non_linearity (str): Type of non-linearity to use.
     """
-
     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):
@@ -346,27 +268,19 @@ class WanMidBlock(nn.Module):
             resnets.append(WanResidualBlock(dim, dim, dropout, non_linearity))
         self.attentions = nn.ModuleList(attentions)
         self.resnets = nn.ModuleList(resnets)
-
         self.gradient_checkpointing = False
 
     def forward(self, x, feat_cache=None, feat_idx=[0]):
-        # First residual block
         x = self.resnets[0](x, feat_cache, feat_idx)
-
-        # Process through attention and residual blocks
         for attn, resnet in zip(self.attentions, self.resnets[1:]):
             if attn is not None:
                 x = attn(x)
-
             x = resnet(x, feat_cache, feat_idx)
-
         return x
 
-
 class WanEncoder3d(nn.Module):
     r"""
     A 3D encoder module.
-
     Args:
         dim (int): The base number of channels in the first layer.
         z_dim (int): The dimensionality of the latent space.
@@ -377,7 +291,6 @@ class WanEncoder3d(nn.Module):
         dropout (float): Dropout rate for the dropout layers.
         non_linearity (str): Type of non-linearity to use.
     """
-
     def __init__(
         self,
         dim=128,
@@ -397,37 +310,23 @@ class WanEncoder3d(nn.Module):
         self.attn_scales = attn_scales
         self.temperal_downsample = temperal_downsample
         self.nonlinearity = get_activation(non_linearity)
-
-        # dimensions
         dims = [dim * u for u in [1] + dim_mult]
         scale = 1.0
-
-        # init block
         self.conv_in = WanCausalConv3d(3, 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:])):
-            # residual (+attention) blocks
             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 forward(self, x, feat_cache=None, feat_idx=[0]):
@@ -435,32 +334,24 @@ class WanEncoder3d(nn.Module):
             idx = feat_idx[0]
             cache_x = x[:, :, -CACHE_T:, :, :].clone()
             if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
-                # cache last frame of last two chunk
                 cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
             x = self.conv_in(x, feat_cache[idx])
             feat_cache[idx] = cache_x
             feat_idx[0] += 1
         else:
             x = self.conv_in(x)
-
-        ## downsamples
         for layer in self.down_blocks:
             if feat_cache is not None:
                 x = layer(x, feat_cache, feat_idx)
             else:
                 x = layer(x)
-
-        ## middle
         x = self.mid_block(x, feat_cache, feat_idx)
-
-        ## head
         x = self.norm_out(x)
         x = self.nonlinearity(x)
         if feat_cache is not None:
             idx = feat_idx[0]
             cache_x = x[:, :, -CACHE_T:, :, :].clone()
             if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
-                # cache last frame of last two chunk
                 cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
             x = self.conv_out(x, feat_cache[idx])
             feat_cache[idx] = cache_x
@@ -469,11 +360,9 @@ class WanEncoder3d(nn.Module):
             x = self.conv_out(x)
         return x
 
-
 class WanUpBlock(nn.Module):
     """
     A block that handles upsampling for the WanVAE decoder.
-
     Args:
         in_dim (int): Input dimension
         out_dim (int): Output dimension
@@ -482,7 +371,6 @@ class WanUpBlock(nn.Module):
         upsample_mode (str, optional): Mode for upsampling ('upsample2d' or 'upsample3d')
         non_linearity (str): Type of non-linearity to use
     """
-
     def __init__(
         self,
         in_dim: int,
@@ -495,42 +383,23 @@ class WanUpBlock(nn.Module):
         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
 
     def forward(self, x, feat_cache=None, feat_idx=[0]):
-        """
-        Forward pass through the upsampling block.
-
-        Args:
-            x (torch.Tensor): Input tensor
-            feat_cache (list, optional): Feature cache for causal convolutions
-            feat_idx (list, optional): Feature index for cache management
-
-        Returns:
-            torch.Tensor: Output tensor
-        """
         for resnet in self.resnets:
             if feat_cache is not None:
                 x = resnet(x, feat_cache, feat_idx)
             else:
                 x = resnet(x)
-
         if self.upsamplers is not None:
             if feat_cache is not None:
                 x = self.upsamplers[0](x, feat_cache, feat_idx)
@@ -538,11 +407,9 @@ class WanUpBlock(nn.Module):
                 x = self.upsamplers[0](x)
         return x
 
-
 class WanDecoder3d(nn.Module):
     r"""
     A 3D decoder module.
-
     Args:
         dim (int): The base number of channels in the first layer.
         z_dim (int): The dimensionality of the latent space.
@@ -553,7 +420,6 @@ class WanDecoder3d(nn.Module):
         dropout (float): Dropout rate for the dropout layers.
         non_linearity (str): Type of non-linearity to use.
     """
-
     def __init__(
         self,
         dim=128,
@@ -572,32 +438,18 @@ class WanDecoder3d(nn.Module):
         self.num_res_blocks = num_res_blocks
         self.attn_scales = attn_scales
         self.temperal_upsample = temperal_upsample
-
         self.nonlinearity = get_activation(non_linearity)
-
-        # dimensions
         dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
         scale = 1.0 / 2 ** (len(dim_mult) - 2)
-
-        # 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:])):
-            # residual (+attention) blocks
             if i > 0:
                 in_dim = in_dim // 2
-
-            # Determine if we need upsampling
             upsample_mode = None
             if i != len(dim_mult) - 1:
                 upsample_mode = "upsample3d" if temperal_upsample[i] else "upsample2d"
-
-            # Create and add the upsampling block
             up_block = WanUpBlock(
                 in_dim=in_dim,
                 out_dim=out_dim,
@@ -607,46 +459,32 @@ class WanDecoder3d(nn.Module):
                 non_linearity=non_linearity,
             )
             self.up_blocks.append(up_block)
-
-            # Update scale for next iteration
             if upsample_mode is not None:
                 scale *= 2.0
-
-        # output blocks
         self.norm_out = WanRMS_norm(out_dim, images=False)
         self.conv_out = WanCausalConv3d(out_dim, 3, 3, padding=1)
-
         self.gradient_checkpointing = False
 
     def forward(self, x, feat_cache=None, feat_idx=[0]):
-        ## conv1
         if feat_cache is not None:
             idx = feat_idx[0]
             cache_x = x[:, :, -CACHE_T:, :, :].clone()
             if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
-                # cache last frame of last two chunk
                 cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
             x = self.conv_in(x, feat_cache[idx])
             feat_cache[idx] = cache_x
             feat_idx[0] += 1
         else:
             x = self.conv_in(x)
-
-        ## middle
         x = self.mid_block(x, feat_cache, feat_idx)
-
-        ## upsamples
         for up_block in self.up_blocks:
             x = up_block(x, feat_cache, feat_idx)
-
-        ## head
         x = self.norm_out(x)
         x = self.nonlinearity(x)
         if feat_cache is not None:
             idx = feat_idx[0]
             cache_x = x[:, :, -CACHE_T:, :, :].clone()
             if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
-                # cache last frame of last two chunk
                 cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
             x = self.conv_out(x, feat_cache[idx])
             feat_cache[idx] = cache_x
@@ -655,16 +493,13 @@ class WanDecoder3d(nn.Module):
             x = self.conv_out(x)
         return x
 
-
 class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
     r"""
     A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos.
     Introduced in [Wan 2.1].
-
     This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
     for all models (such as downloading or saving).
     """
-
     _supports_gradient_checkpointing = False
 
     @register_to_config
@@ -678,54 +513,23 @@ class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
         temperal_downsample: List[bool] = [False, True, True],
         dropout: float = 0.0,
         latents_mean: List[float] = [
-            -0.7571,
-            -0.7089,
-            -0.9113,
-            0.1075,
-            -0.1745,
-            0.9653,
-            -0.1517,
-            1.5508,
-            0.4134,
-            -0.0715,
-            0.5517,
-            -0.3632,
-            -0.1922,
-            -0.9497,
-            0.2503,
-            -0.2921,
+            -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,
+            0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921,
         ],
         latents_std: List[float] = [
-            2.8184,
-            1.4541,
-            2.3275,
-            2.6558,
-            1.2196,
-            1.7708,
-            2.6052,
-            2.0743,
-            3.2687,
-            2.1526,
-            2.8652,
-            1.5579,
-            1.6382,
-            1.1253,
-            2.8251,
-            1.9160,
+            2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,
+            3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160,
         ],
     ) -> None:
         super().__init__()
-
         self.z_dim = z_dim
         self.temperal_downsample = temperal_downsample
         self.temperal_upsample = temperal_downsample[::-1]
-
         self.encoder = WanEncoder3d(
             base_dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout
         )
         self.quant_conv = WanCausalConv3d(z_dim * 2, z_dim * 2, 1)
         self.post_quant_conv = WanCausalConv3d(z_dim, z_dim, 1)
-
         self.decoder = WanDecoder3d(
             base_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout
         )
@@ -741,14 +545,12 @@ class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
         self._conv_num = _count_conv3d(self.decoder)
         self._conv_idx = [0]
         self._feat_map = [None] * self._conv_num
-        # cache encode
         self._enc_conv_num = _count_conv3d(self.encoder)
         self._enc_conv_idx = [0]
         self._enc_feat_map = [None] * self._enc_conv_num
 
     def _encode(self, x: torch.Tensor) -> torch.Tensor:
         self.clear_cache()
-        ## cache
         t = x.shape[2]
         iter_ = 1 + (t - 1) // 4
         for i in range(iter_):
@@ -762,7 +564,6 @@ class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
                     feat_idx=self._enc_conv_idx,
                 )
                 out = torch.cat([out, out_], 2)
-
         enc = self.quant_conv(out)
         mu, logvar = enc[:, : self.z_dim, :, :, :], enc[:, self.z_dim :, :, :, :]
         enc = torch.cat([mu, logvar], dim=1)
@@ -775,12 +576,10 @@ class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
     ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
         r"""
         Encode a batch of images into latents.
-
         Args:
             x (`torch.Tensor`): Input batch of images.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
-
         Returns:
                 The latent representations of the encoded videos. If `return_dict` is True, a
                 [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
@@ -793,7 +592,6 @@ class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
 
     def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
         self.clear_cache()
-
         iter_ = z.shape[2]
         x = self.post_quant_conv(z)
         for i in range(iter_):
@@ -803,24 +601,20 @@ class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
             else:
                 out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)
                 out = torch.cat([out, out_], 2)
-
         out = torch.clamp(out, min=-1.0, max=1.0)
         self.clear_cache()
         if not return_dict:
             return (out,)
-
         return DecoderOutput(sample=out)
 
     @apply_forward_hook
     def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
         r"""
         Decode a batch of images.
-
         Args:
             z (`torch.Tensor`): Input batch of latent vectors.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
-
         Returns:
             [`~models.vae.DecoderOutput`] or `tuple`:
                 If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
@@ -829,7 +623,6 @@ class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
         decoded = self._decode(z).sample
         if not return_dict:
             return (decoded,)
-
         return DecoderOutput(sample=decoded)
 
     def forward(
@@ -852,4 +645,4 @@ class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
         else:
             z = posterior.mode()
         dec = self.decode(z, return_dict=return_dict)
-        return dec
+        return dec