Add files via upload

nodes.py

@@ -35,7 +35,7 @@ def add_noise_to_reference_video(image, ratio=None):
     image = image + image_noise
     return image
 
-class WanVideoBlockSwap:
+class WanVideoBlockSwap_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -54,7 +54,7 @@ class WanVideoBlockSwap:
     def setargs(self, **kwargs):
         return (kwargs, )
 
-class WanVideoVRAMManagement:
+class WanVideoVRAMManagement_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -71,7 +71,7 @@ class WanVideoVRAMManagement:
     def setargs(self, **kwargs):
         return (kwargs, )
 
-class WanVideoTeaCache:
+class WanVideoTeaCache_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -127,7 +127,7 @@ Official recommended values https://github.com/ali-vilab/TeaCache/tree/main/TeaC
         return (teacache_args,)
 
 
-class WanVideoModel(comfy.model_base.BaseModel):
+class WanVideoModel_2frames(comfy.model_base.BaseModel):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.pipeline = {}
@@ -146,7 +146,7 @@ except: #for backwards compatibility
     from comfy.latent_formats import HunyuanVideo
     latent_format = HunyuanVideo
 
-class WanVideoModelConfig:
+class WanVideoModelConfig_2frames:
     def __init__(self, dtype):
         self.unet_config = {}
         self.unet_extra_config = {}
@@ -204,7 +204,7 @@ def standardize_lora_key_format(lora_sd):
         new_sd[k] = v
     return new_sd
 
-class WanVideoEnhanceAVideo:
+class WanVideoEnhanceAVideo_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -223,7 +223,7 @@ class WanVideoEnhanceAVideo:
     def setargs(self, **kwargs):
         return (kwargs, )
 
-class WanVideoLoraSelect:
+class WanVideoLoraSelect_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -259,7 +259,7 @@ class WanVideoLoraSelect:
         loras_list.append(lora)
         return (loras_list,)
 
-class WanVideoLoraBlockEdit:
+class WanVideoLoraBlockEdit_2frames:
     def __init__(self):
         self.loaded_lora = None
 
@@ -286,7 +286,7 @@ class WanVideoLoraBlockEdit:
         return (selected_blocks,)
 
 #region Model loading
-class WanVideoModelLoader:
+class WanVideoModelLoader_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -407,8 +407,8 @@ class WanVideoModelLoader:
             transformer = WanModel(**TRANSFORMER_CONFIG)
         transformer.eval()
 
-        comfy_model = WanVideoModel(
-            WanVideoModelConfig(base_dtype),
+        comfy_model = WanVideoModel_2frames(
+            WanVideoModelConfig_2frames(base_dtype),
             model_type=comfy.model_base.ModelType.FLOW,
             device=device,
         )
@@ -586,7 +586,7 @@ class WanVideoModelLoader:
 
 #region load VAE
 
-class WanVideoVAELoader:
+class WanVideoVAELoader_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -607,7 +607,7 @@ class WanVideoVAELoader:
     DESCRIPTION = "Loads Hunyuan VAE model from 'ComfyUI/models/vae'"
 
     def loadmodel(self, model_name, precision):
-        from .wanvideo.wan_video_vae import WanVideoVAE
+        from .wanvideo.wan_video_vae_2frames import WanVideoVAE
 
         device = mm.get_torch_device()
         offload_device = mm.unet_offload_device()
@@ -632,7 +632,7 @@ class WanVideoVAELoader:
 
 
 
-class WanVideoTorchCompileSettings:
+class WanVideoTorchCompileSettings_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -667,7 +667,7 @@ class WanVideoTorchCompileSettings:
 
 #region TextEncode
 
-class LoadWanVideoT5TextEncoder:
+class LoadWanVideoT5TextEncoder_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -718,7 +718,7 @@ class LoadWanVideoT5TextEncoder:
         
         return (text_encoder,)
     
-class LoadWanVideoClipTextEncoder:
+class LoadWanVideoClipTextEncoder_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -757,7 +757,7 @@ class LoadWanVideoClipTextEncoder:
         return (clip_model,)
     
 
-class WanVideoTextEncode:
+class WanVideoTextEncode_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
@@ -806,7 +806,7 @@ class WanVideoTextEncode:
             }
         return (prompt_embeds_dict,)
     
-class WanVideoTextEmbedBridge:
+class WanVideoTextEmbedBridge_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
@@ -830,7 +830,7 @@ class WanVideoTextEmbedBridge:
         return (prompt_embeds_dict,)
     
 #region clip image encode
-class WanVideoImageClipEncode:
+class WanVideoImageClipEncode_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
@@ -959,7 +959,7 @@ class WanVideoImageClipEncode:
         return (image_embeds,)
 
 
-class WanVideoStartEndImageClipEncode:
+class WanVideoStartEndImageClipEncode_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
@@ -1102,7 +1102,7 @@ class WanVideoStartEndImageClipEncode:
 
 
 
-class WanVideoEmptyEmbeds:
+class WanVideoEmptyEmbeds_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
@@ -1141,7 +1141,7 @@ class WanVideoEmptyEmbeds:
 
 #region Sampler
 
-class WanVideoContextOptions:
+class WanVideoContextOptions_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
@@ -1172,7 +1172,7 @@ class WanVideoContextOptions:
 
         return (context_options,)
     
-class WanVideoFlowEdit:
+class WanVideoFlowEdit_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
@@ -1197,7 +1197,7 @@ class WanVideoFlowEdit:
     def process(self, **kwargs):
         return (kwargs,)
     
-class WanVideoSampler:
+class WanVideoSampler_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -1347,7 +1347,7 @@ class WanVideoSampler:
             context_stride = context_options["context_stride"] // 4
             context_overlap = context_options["context_overlap"] // 4
 
-            self.window_tracker = WindowTracker(verbose=context_options["verbose"])
+            self.window_tracker = WindowTracker_2frames(verbose=context_options["verbose"])
 
             # Get total number of prompts
             num_prompts = len(text_embeds["prompt_embeds"])
@@ -1821,7 +1821,7 @@ class WanVideoSampler:
             }, )
 
 
-class WanVideoStartEndSampler:
+class WanVideoStartEndSampler_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -1972,7 +1972,7 @@ class WanVideoStartEndSampler:
             context_stride = context_options["context_stride"] // 4
             context_overlap = context_options["context_overlap"] // 4
 
-            self.window_tracker = WindowTracker(verbose=context_options["verbose"])
+            self.window_tracker = WindowTracker_2frames(verbose=context_options["verbose"])
 
             # Get total number of prompts
             num_prompts = len(text_embeds["prompt_embeds"])
@@ -2460,7 +2460,7 @@ class WanVideoStartEndSampler:
         return ({
                     "samples": x0.unsqueeze(0).cpu()
                 },)
-class WindowTracker:
+class WindowTracker_2frames:
     def __init__(self, verbose=False):
         self.window_map = {}  # Maps frame sequence to persistent ID
         self.next_id = 0
@@ -2484,7 +2484,7 @@ class WindowTracker:
         return self.teacache_states[window_id]
 
 #region VideoDecode
-class WanVideoDecode:
+class WanVideoDecode_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
@@ -2527,7 +2527,7 @@ class WanVideoDecode:
 
         return (image,)
 
-class WanVideoStartEndDecode:
+class WanVideoStartEndDecode_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
@@ -2580,7 +2580,7 @@ class WanVideoStartEndDecode:
 
 
 #region VideoEncode
-class WanVideoEncode:
+class WanVideoEncode_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
@@ -2650,7 +2650,7 @@ class WanVideoEncode:
 
         return ({"samples": latents, "mask": mask},)
 
-class WanVideoLatentPreview:
+class WanVideoLatentPreview_2frames:
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -2730,50 +2730,50 @@ class WanVideoLatentPreview:
         return (latent_images.float().cpu(), out_factors)
 
 NODE_CLASS_MAPPINGS = {
-    "WanVideoTextEncode": WanVideoTextEncode,
-    "WanVideoModelLoader": WanVideoModelLoader,
-    "WanVideoVAELoader": WanVideoVAELoader,
-    "LoadWanVideoT5TextEncoder": LoadWanVideoT5TextEncoder,
-    "WanVideoImageClipEncode": WanVideoImageClipEncode,
-    "LoadWanVideoClipTextEncoder": LoadWanVideoClipTextEncoder,
-    "WanVideoEncode": WanVideoEncode,
-    "WanVideoBlockSwap": WanVideoBlockSwap,
-    "WanVideoTorchCompileSettings": WanVideoTorchCompileSettings,
-    "WanVideoLatentPreview": WanVideoLatentPreview,
-    "WanVideoEmptyEmbeds": WanVideoEmptyEmbeds,
-    "WanVideoLoraSelect": WanVideoLoraSelect,
-    "WanVideoLoraBlockEdit": WanVideoLoraBlockEdit,
-    "WanVideoEnhanceAVideo": WanVideoEnhanceAVideo,
-    "WanVideoContextOptions": WanVideoContextOptions,
-    "WanVideoTeaCache": WanVideoTeaCache,
-    "WanVideoVRAMManagement": WanVideoVRAMManagement,
-    "WanVideoTextEmbedBridge": WanVideoTextEmbedBridge,
-    "WanVideoFlowEdit": WanVideoFlowEdit,
-    "WanVideoStartEndImageClipEncode":WanVideoStartEndImageClipEncode,
-    "WanVideoStartEndDecode": WanVideoStartEndDecode,
-    "WanVideoStartEndSampler":WanVideoStartEndSampler,
+    "WanVideoTextEncode_2frames": WanVideoTextEncode_2frames,
+    "WanVideoModelLoader_2frames": WanVideoModelLoader_2frames,
+    "WanVideoVAELoader_2frames": WanVideoVAELoader_2frames,
+    "LoadWanVideoT5TextEncoder_2frames": LoadWanVideoT5TextEncoder_2frames,
+    "WanVideoImageClipEncode_2frames": WanVideoImageClipEncode_2frames,
+    "LoadWanVideoClipTextEncoder_2frames": LoadWanVideoClipTextEncoder_2frames,
+    "WanVideoEncode_2frames": WanVideoEncode_2frames,
+    "WanVideoBlockSwap_2frames": WanVideoBlockSwap_2frames,
+    "WanVideoTorchCompileSettings_2frames": WanVideoTorchCompileSettings_2frames,
+    "WanVideoLatentPreview_2frames": WanVideoLatentPreview_2frames,
+    "WanVideoEmptyEmbeds_2frames": WanVideoEmptyEmbeds_2frames,
+    "WanVideoLoraSelect_2frames": WanVideoLoraSelect_2frames,
+    "WanVideoLoraBlockEdit_2frames": WanVideoLoraBlockEdit_2frames,
+    "WanVideoEnhanceAVideo_2frames": WanVideoEnhanceAVideo_2frames,
+    "WanVideoContextOptions_2frames": WanVideoContextOptions_2frames,
+    "WanVideoTeaCache_2frames": WanVideoTeaCache_2frames,
+    "WanVideoVRAMManagement_2frames": WanVideoVRAMManagement_2frames,
+    "WanVideoTextEmbedBridge_2frames": WanVideoTextEmbedBridge_2frames,
+    "WanVideoFlowEdit_2frames": WanVideoFlowEdit_2frames,
+    "WanVideoStartEndImageClipEncode_2frames":WanVideoStartEndImageClipEncode_2frames,
+    "WanVideoStartEndDecode_2frames": WanVideoStartEndDecode_2frames,
+    "WanVideoStartEndSampler_2frames":WanVideoStartEndSampler_2frames,
     }
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "WanVideoTextEncode": "WanVideo TextEncode(2Frames)",
-    "WanVideoModelLoader": "WanVideo Model Loader(2Frames)",
-    "WanVideoVAELoader": "WanVideo VAE Loader(2Frames)",
-    "LoadWanVideoT5TextEncoder": "Load WanVideo T5 TextEncoder(2Frames)",
-    "WanVideoImageClipEncode": "WanVideo ImageClip Encode(2Frames)",
-    "LoadWanVideoClipTextEncoder": "Load WanVideo Clip Encoder(2Frames)",
-    "WanVideoEncode": "WanVideo Encode(2Frames)",
-    "WanVideoBlockSwap": "WanVideo BlockSwap(2Frames)",
-    "WanVideoTorchCompileSettings": "WanVideo Torch Compile Settings(2Frames)",
-    "WanVideoLatentPreview": "WanVideo Latent Preview(2Frames)",
-    "WanVideoEmptyEmbeds": "WanVideo Empty Embeds(2Frames)",
-    "WanVideoLoraSelect": "WanVideo Lora Select(2Frames)",
-    "WanVideoLoraBlockEdit": "WanVideo Lora Block Edit(2Frames)",
-    "WanVideoEnhanceAVideo": "WanVideo Enhance-A-Video(2Frames)",
-    "WanVideoContextOptions": "WanVideo Context Options(2Frames)",
-    "WanVideoTeaCache": "WanVideo TeaCache(2Frames)",
-    "WanVideoVRAMManagement": "WanVideo VRAM Management(2Frames)",
-    "WanVideoTextEmbedBridge": "WanVideo TextEmbed Bridge(2Frames)",
-    "WanVideoFlowEdit": "WanVideo FlowEdit(2Frames)",
-    "WanVideoStartEndImageClipEncode":"WanVideo Start End Image Clip Encode(2Frames)",
-    "WanVideoStartEndDecode": "Wan Video Start End Decode(2Frames)",
-    "WanVideoStartEndSampler": "Wan Video Start End Sampler(2Frames)",
+    "WanVideoTextEncode_2frames": "WanVideo TextEncode(2Frames)",
+    "WanVideoModelLoader_2frames": "WanVideo Model Loader(2Frames)",
+    "WanVideoVAELoader_2frames": "WanVideo VAE Loader(2Frames)",
+    "LoadWanVideoT5TextEncoder_2frames": "Load WanVideo T5 TextEncoder(2Frames)",
+    "WanVideoImageClipEncode_2frames": "WanVideo ImageClip Encode(2Frames)",
+    "LoadWanVideoClipTextEncoder_2frames": "Load WanVideo Clip Encoder(2Frames)",
+    "WanVideoEncode_2frames": "WanVideo Encode(2Frames)",
+    "WanVideoBlockSwap_2frames": "WanVideo BlockSwap(2Frames)",
+    "WanVideoTorchCompileSettings_2frames": "WanVideo Torch Compile Settings(2Frames)",
+    "WanVideoLatentPreview_2frames": "WanVideo Latent Preview(2Frames)",
+    "WanVideoEmptyEmbeds_2frames": "WanVideo Empty Embeds(2Frames)",
+    "WanVideoLoraSelect_2frames": "WanVideo Lora Select(2Frames)",
+    "WanVideoLoraBlockEdit_2frames": "WanVideo Lora Block Edit(2Frames)",
+    "WanVideoEnhanceAVideo_2frames": "WanVideo Enhance-A-Video(2Frames)",
+    "WanVideoContextOptions_2frames": "WanVideo Context Options(2Frames)",
+    "WanVideoTeaCache_2frames": "WanVideo TeaCache(2Frames)",
+    "WanVideoVRAMManagement_2frames": "WanVideo VRAM Management(2Frames)",
+    "WanVideoTextEmbedBridge_2frames": "WanVideo TextEmbed Bridge(2Frames)",
+    "WanVideoFlowEdit_2frames": "WanVideo FlowEdit(2Frames)",
+    "WanVideoStartEndImageClipEncode_2frames":"WanVideo Start End Image Clip Encode(2Frames)",
+    "WanVideoStartEndDecode_2frames": "Wan Video Start End Decode(2Frames)",
+    "WanVideoStartEndSampler_2frames": "Wan Video Start End Sampler(2Frames)",
     }

wanvideo/wan_video_vae_2frames.py

@@ -0,0 +1,903 @@
+from einops import rearrange, repeat
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+
+CACHE_T = 2
+
+
+def check_is_instance(model, module_class):
+    if isinstance(model, module_class):
+        return True
+    if hasattr(model, "module") and isinstance(model.module, module_class):
+        return True
+    return False
+
+
+def block_causal_mask(x, block_size):
+    # params
+    b, n, s, _, device = *x.size(), x.device
+    assert s % block_size == 0
+    num_blocks = s // block_size
+
+    # build mask
+    mask = torch.zeros(b, n, s, s, dtype=torch.bool, device=device)
+    for i in range(num_blocks):
+        mask[:, :,
+             i * block_size:(i + 1) * block_size, :(i + 1) * block_size] = 1
+    return mask
+
+
+class CausalConv3d(nn.Conv3d):
+    """
+    Causal 3d convolusion.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._padding = (self.padding[2], self.padding[2], self.padding[1],
+                         self.padding[1], 2 * self.padding[0], 0)
+        self.padding = (0, 0, 0)
+
+    def forward(self, x, cache_x=None):
+        padding = list(self._padding)
+        if cache_x is not None and self._padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+        x = F.pad(x, padding)
+
+        return super().forward(x)
+
+
+class RMS_norm(nn.Module):
+
+    def __init__(self, dim, channel_first=True, images=True, bias=False):
+        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.
+
+    def forward(self, x):
+        return F.normalize(
+            x, dim=(1 if self.channel_first else
+                    -1)) * self.scale * self.gamma + self.bias
+
+
+class Upsample(nn.Upsample):
+
+    def forward(self, x):
+        """
+        Fix bfloat16 support for nearest neighbor interpolation.
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+
+    def __init__(self, dim, mode):
+        assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d',
+                        'downsample3d')
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        # layers
+        if mode == 'upsample2d':
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2., 2.), mode='nearest-exact'),
+                nn.Conv2d(dim, dim // 2, 3, padding=1))
+        elif mode == 'upsample3d':
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2., 2.), mode='nearest-exact'),
+                nn.Conv2d(dim, dim // 2, 3, padding=1))
+            self.time_conv = CausalConv3d(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 = CausalConv3d(dim,
+                                          dim, (3, 1, 1),
+                                          stride=(2, 1, 1),
+                                          padding=(0, 0, 0))
+
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        b, c, t, h, w = x.size()
+        if self.mode == 'upsample3d':
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = 'Rep'
+                    feat_idx[0] += 1
+                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)
+                    if cache_x.shape[2] < 2 and feat_cache[
+                            idx] is not None and feat_cache[idx] == 'Rep':
+                        cache_x = torch.cat([
+                            torch.zeros_like(cache_x).to(cache_x.device),
+                            cache_x
+                        ],
+                                            dim=2)
+                    if feat_cache[idx] == 'Rep':
+                        x = self.time_conv(x)
+                    else:
+                        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)
+        t = x.shape[2]
+        x = rearrange(x, 'b c t h w -> (b t) c h w')
+        x = self.resample(x)
+        x = rearrange(x, '(b t) c h w -> b c t h w', t=t)
+
+        if self.mode == 'downsample3d':
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = x.clone()
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -1:, :, :].clone()
+                    x = self.time_conv(
+                        torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+        return x
+
+    def init_weight(self, conv):
+        conv_weight = conv.weight
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        one_matrix = torch.eye(c1, c2)
+        init_matrix = one_matrix
+        nn.init.zeros_(conv_weight)
+        conv_weight.data[:, :, 1, 0, 0] = init_matrix
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+    def init_weight2(self, conv):
+        conv_weight = conv.weight.data
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        init_matrix = torch.eye(c1 // 2, c2)
+        conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix
+        conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+
+class ResidualBlock(nn.Module):
+
+    def __init__(self, in_dim, out_dim, dropout=0.0):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+        # layers
+        self.residual = nn.Sequential(
+            RMS_norm(in_dim, images=False), nn.SiLU(),
+            CausalConv3d(in_dim, out_dim, 3, padding=1),
+            RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout),
+            CausalConv3d(out_dim, out_dim, 3, padding=1))
+        self.shortcut = CausalConv3d(in_dim, out_dim, 1) \
+            if in_dim != out_dim else nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        h = self.shortcut(x)
+        for layer in self.residual:
+            if check_is_instance(layer, CausalConv3d) and 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 = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x + h
+
+
+class AttentionBlock(nn.Module):
+    """
+    Causal self-attention with a single head.
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+        # layers
+        self.norm = RMS_norm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+
+        # zero out the last layer params
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, x):
+        identity = x
+        b, c, t, h, w = x.size()
+        x = rearrange(x, 'b c t h w -> (b t) c h w')
+        x = self.norm(x)
+        # compute query, key, value
+        q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(
+            0, 1, 3, 2).contiguous().chunk(3, dim=-1)
+
+        # apply attention
+        x = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            #attn_mask=block_causal_mask(q, block_size=h * w)
+        )
+        x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+
+        # output
+        x = self.proj(x)
+        x = rearrange(x, '(b t) c h w-> b c t h w', t=t)
+        return x + identity
+
+
+class Encoder3d(nn.Module):
+
+    def __init__(self,
+                 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):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+
+        # dimensions
+        dims = [dim * u for u in [1] + dim_mult]
+        scale = 1.0
+
+        # init block
+        self.conv1 = CausalConv3d(3, dims[0], 3, padding=1)
+
+        # downsample blocks
+        downsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            # residual (+attention) blocks
+            for _ in range(num_res_blocks):
+                downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    downsamples.append(AttentionBlock(out_dim))
+                in_dim = out_dim
+
+            # downsample block
+            if i != len(dim_mult) - 1:
+                mode = 'downsample3d' if temperal_downsample[
+                    i] else 'downsample2d'
+                downsamples.append(Resample(out_dim, mode=mode))
+                scale /= 2.0
+        self.downsamples = nn.Sequential(*downsamples)
+
+        # middle blocks
+        self.middle = nn.Sequential(ResidualBlock(out_dim, out_dim, dropout),
+                                    AttentionBlock(out_dim),
+                                    ResidualBlock(out_dim, out_dim, dropout))
+
+        # output blocks
+        self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(),
+                                  CausalConv3d(out_dim, z_dim, 3, padding=1))
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        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.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        ## downsamples
+        for layer in self.downsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## middle
+        for layer in self.middle:
+            if check_is_instance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## head
+        for layer in self.head:
+            if check_is_instance(layer, CausalConv3d) and 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 = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+class Decoder3d(nn.Module):
+
+    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):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_upsample = temperal_upsample
+
+        # dimensions
+        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        scale = 1.0 / 2**(len(dim_mult) - 2)
+
+        # init block
+        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+
+        # middle blocks
+        self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout),
+                                    AttentionBlock(dims[0]),
+                                    ResidualBlock(dims[0], dims[0], dropout))
+
+        # upsample blocks
+        upsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            # residual (+attention) blocks
+            if i == 1 or i == 2 or i == 3:
+                in_dim = in_dim // 2
+            for _ in range(num_res_blocks + 1):
+                upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    upsamples.append(AttentionBlock(out_dim))
+                in_dim = out_dim
+
+            # upsample block
+            if i != len(dim_mult) - 1:
+                mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d'
+                upsamples.append(Resample(out_dim, mode=mode))
+                scale *= 2.0
+        self.upsamples = nn.Sequential(*upsamples)
+
+        # output blocks
+        self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(),
+                                  CausalConv3d(out_dim, 3, 3, padding=1))
+
+    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.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        ## middle
+        for layer in self.middle:
+            if check_is_instance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## upsamples
+        for layer in self.upsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## head
+        for layer in self.head:
+            if check_is_instance(layer, CausalConv3d) and 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 = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+def count_conv3d(model):
+    count = 0
+    for m in model.modules():
+        if check_is_instance(m, CausalConv3d):
+            count += 1
+    return count
+
+
+class VideoVAE_(nn.Module):
+
+    def __init__(self,
+                 dim=96,
+                 z_dim=16,
+                 dim_mult=[1, 2, 4, 4],
+                 num_res_blocks=2,
+                 attn_scales=[],
+                 temperal_downsample=[False, True, True],
+                 dropout=0.0,):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+        self.temperal_upsample = temperal_downsample[::-1]
+
+        # modules
+        self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks,
+                                 attn_scales, self.temperal_downsample, dropout)
+        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+        self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks,
+                                 attn_scales, self.temperal_upsample, dropout)
+
+    def forward(self, x):
+        mu, log_var = self.encode(x)
+        z = self.reparameterize(mu, log_var)
+        x_recon = self.decode(z)
+        return x_recon, mu, log_var
+
+
+    def encode_2(self, x, scale):
+        self.clear_cache()
+        ## cache
+        t = x.shape[2]
+        iter_ = 2 + (t - 2) // 4
+
+        for i in range(iter_):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self.encoder(x[:, :, :1, :, :],
+                                   feat_cache=self._enc_feat_map,
+                                   feat_idx=self._enc_conv_idx)
+            elif i== iter_-1:
+                out_ = self.encoder(x[:, :, -1:, :, :],
+                                   feat_cache=[None] * self._enc_conv_num,
+                                   feat_idx=self._enc_conv_idx)
+                out = torch.cat([out, out_], 2)
+            else:
+                out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
+                                    feat_cache=self._enc_feat_map,
+                                    feat_idx=self._enc_conv_idx)
+                out = torch.cat([out, out_], 2)
+        out_head = out[:, :, :iter_ - 1, :, :]
+        out_tail = out[:, :, -1, :, :].unsqueeze(2)
+        mu, log_var = torch.cat([self.conv1(out_head), self.conv1(out_tail)], dim=2).chunk(2, dim=1)
+        if isinstance(scale[0], torch.Tensor):
+            scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale]
+            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            scale = scale.to(dtype=mu.dtype, device=mu.device)
+            mu = (mu - scale[0]) * scale[1]
+        return mu
+
+
+    def encode(self, x, scale):
+        self.clear_cache()
+        ## cache
+        t = x.shape[2]
+        iter_ = 1 + (t - 1) // 4
+
+        for i in range(iter_):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self.encoder(x[:, :, :1, :, :],
+                                   feat_cache=self._enc_feat_map,
+                                   feat_idx=self._enc_conv_idx)
+            else:
+                out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
+                                    feat_cache=self._enc_feat_map,
+                                    feat_idx=self._enc_conv_idx)
+                out = torch.cat([out, out_], 2)
+        mu, log_var = self.conv1(out).chunk(2, dim=1)
+        if isinstance(scale[0], torch.Tensor):
+            scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale]
+            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            scale = scale.to(dtype=mu.dtype, device=mu.device)
+            mu = (mu - scale[0]) * scale[1]
+        return mu
+
+
+
+    def decode_2(self, z, scale):
+        print("decode_2222222222")
+        self.clear_cache()
+        # z: [b,c,t,h,w]
+        if isinstance(scale[0], torch.Tensor):
+            scale = [s.to(dtype=z.dtype, device=z.device) for s in scale]
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            scale = scale.to(dtype=z.dtype, device=z.device)
+            z = z / scale[1] + scale[0]
+        iter_ = z.shape[2]
+        z_head=z[:,:,:-1,:,:]
+        z_tail=z[:,:,-1,:,:].unsqueeze(2)
+        x = torch.cat([self.conv2(z_head), self.conv2(z_tail)], dim=2)
+        for i in range(iter_):
+            self._conv_idx = [0]
+            if i == 0:
+                out = self.decoder(x[:, :, i:i + 1, :, :],
+                                   feat_cache=self._feat_map,
+                                   feat_idx=self._conv_idx)
+
+            elif i==iter_-1:
+                out_ = self.decoder(x[:, :, -1, :, :],
+                                    feat_cache=[None] * self._enc_conv_num,
+                                    feat_idx=self._conv_idx)
+                out = torch.cat([out, out_], 2) # may add tensor offload
+            else:
+                out_ = self.decoder(x[:, :, i:i + 1, :, :],
+                                    feat_cache=self._feat_map,
+                                    feat_idx=self._conv_idx)
+                out = torch.cat([out, out_], 2) # may add tensor offload
+            print(out.shape,'out')
+        return out
+
+
+
+    def decode(self, z, scale):
+        self.clear_cache()
+        # z: [b,c,t,h,w]
+        if isinstance(scale[0], torch.Tensor):
+            scale = [s.to(dtype=z.dtype, device=z.device) for s in scale]
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            scale = scale.to(dtype=z.dtype, device=z.device)
+            z = z / scale[1] + scale[0]
+        iter_ = z.shape[2]
+        x = self.conv2(z)
+        for i in range(iter_):
+            self._conv_idx = [0]
+            if i == 0:
+                out = self.decoder(x[:, :, i:i + 1, :, :],
+                                   feat_cache=self._feat_map,
+                                   feat_idx=self._conv_idx)
+            else:
+                out_ = self.decoder(x[:, :, i:i + 1, :, :],
+                                    feat_cache=self._feat_map,
+                                    feat_idx=self._conv_idx)
+                out = torch.cat([out, out_], 2) # may add tensor offload
+        return out
+
+    def reparameterize(self, mu, log_var):
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return eps * std + mu
+
+    def sample(self, imgs, deterministic=False):
+        mu, log_var = self.encode(imgs)
+        if deterministic:
+            return mu
+        std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+        return mu + std * torch.randn_like(std)
+
+    def clear_cache(self):
+        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
+
+
+class WanVideoVAE(nn.Module):
+
+    def __init__(self, z_dim=16, dtype=torch.float32):
+        super().__init__()
+
+        self.dtype = dtype
+
+        mean = [
+            -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
+        ]
+        std = [
+            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
+        ]
+        self.mean = torch.tensor(mean)
+        self.std = torch.tensor(std)
+        self.scale = [self.mean, 1.0 / self.std]
+
+        # init model
+        self.model = VideoVAE_(z_dim=z_dim).eval().requires_grad_(False)
+        self.upsampling_factor = 8
+
+
+    def build_1d_mask(self, length, left_bound, right_bound, border_width):
+        x = torch.ones((length,))
+        if not left_bound:
+            x[:border_width] = (torch.arange(border_width) + 1) / border_width
+        if not right_bound:
+            x[-border_width:] = torch.flip((torch.arange(border_width) + 1) / border_width, dims=(0,))
+        return x
+
+
+    def build_mask(self, data, is_bound, border_width):
+        _, _, _, H, W = data.shape
+        h = self.build_1d_mask(H, is_bound[0], is_bound[1], border_width[0])
+        w = self.build_1d_mask(W, is_bound[2], is_bound[3], border_width[1])
+
+        h = repeat(h, "H -> H W", H=H, W=W)
+        w = repeat(w, "W -> H W", H=H, W=W)
+
+        mask = torch.stack([h, w]).min(dim=0).values
+        mask = rearrange(mask, "H W -> 1 1 1 H W")
+        return mask
+
+
+    def tiled_decode(self, hidden_states, device, tile_size, tile_stride):
+        _, _, T, H, W = hidden_states.shape
+        size_h, size_w = tile_size
+        stride_h, stride_w = tile_stride
+
+        # Split tasks
+        tasks = []
+        for h in range(0, H, stride_h):
+            if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue
+            for w in range(0, W, stride_w):
+                if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue
+                h_, w_ = h + size_h, w + size_w
+                tasks.append((h, h_, w, w_))
+
+        data_device = "cpu"
+        computation_device = device
+
+        out_T = T * 4 - 3
+        weight = torch.zeros((1, 1, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=data_device)
+        values = torch.zeros((1, 3, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=data_device)
+
+        for h, h_, w, w_ in tqdm(tasks, desc="VAE decoding"):
+            hidden_states_batch = hidden_states[:, :, :, h:h_, w:w_].to(computation_device)
+            hidden_states_batch = self.model.decode(hidden_states_batch, self.scale).to(data_device)
+
+            mask = self.build_mask(
+                hidden_states_batch,
+                is_bound=(h==0, h_>=H, w==0, w_>=W),
+                border_width=((size_h - stride_h) * self.upsampling_factor, (size_w - stride_w) * self.upsampling_factor)
+            ).to(dtype=hidden_states.dtype, device=data_device)
+
+            target_h = h * self.upsampling_factor
+            target_w = w * self.upsampling_factor
+            values[
+                :,
+                :,
+                :,
+                target_h:target_h + hidden_states_batch.shape[3],
+                target_w:target_w + hidden_states_batch.shape[4],
+            ] += hidden_states_batch * mask
+            weight[
+                :,
+                :,
+                :,
+                target_h: target_h + hidden_states_batch.shape[3],
+                target_w: target_w + hidden_states_batch.shape[4],
+            ] += mask
+        values = values / weight
+        values = values.float().clamp_(-1, 1)
+        return values
+
+
+    def tiled_encode(self, video, device, tile_size, tile_stride):
+        _, _, T, H, W = video.shape
+        size_h, size_w = tile_size
+        stride_h, stride_w = tile_stride
+
+        # Split tasks
+        tasks = []
+        for h in range(0, H, stride_h):
+            if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue
+            for w in range(0, W, stride_w):
+                if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue
+                h_, w_ = h + size_h, w + size_w
+                tasks.append((h, h_, w, w_))
+
+        data_device = "cpu"
+        computation_device = device
+
+        out_T = (T + 3) // 4
+        weight = torch.zeros((1, 1, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device)
+        values = torch.zeros((1, 16, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device)
+
+        for h, h_, w, w_ in tqdm(tasks, desc="VAE encoding"):
+            hidden_states_batch = video[:, :, :, h:h_, w:w_].to(computation_device)
+            hidden_states_batch = self.model.encode(hidden_states_batch, self.scale).to(data_device)
+
+            mask = self.build_mask(
+                hidden_states_batch,
+                is_bound=(h==0, h_>=H, w==0, w_>=W),
+                border_width=((size_h - stride_h) // self.upsampling_factor, (size_w - stride_w) // self.upsampling_factor)
+            ).to(dtype=video.dtype, device=data_device)
+
+            target_h = h // self.upsampling_factor
+            target_w = w // self.upsampling_factor
+            values[
+                :,
+                :,
+                :,
+                target_h:target_h + hidden_states_batch.shape[3],
+                target_w:target_w + hidden_states_batch.shape[4],
+            ] += hidden_states_batch * mask
+            weight[
+                :,
+                :,
+                :,
+                target_h: target_h + hidden_states_batch.shape[3],
+                target_w: target_w + hidden_states_batch.shape[4],
+            ] += mask
+        values = values / weight
+        values = values.float()
+        return values
+
+
+    def single_encode(self, video, device):
+        video = video.to(device)
+        x = self.model.encode(video, self.scale)
+        return x.float()
+
+
+    def single_decode(self, hidden_state, device):
+        print("single_decode")
+        hidden_state = hidden_state.to(device)
+        video = self.model.decode(hidden_state, self.scale)
+        return video.float().clamp_(-1, 1)
+
+    def double_encode(self, video, device):
+        print('double_encode')
+        video = video.to(device)
+        x = self.model.encode_2(video, self.scale)
+        return x.float()
+
+
+    def double_decode(self, hidden_state, device):
+        print('double_decode')
+        hidden_state = hidden_state.to(device)
+        video = self.model.decode_2(hidden_state, self.scale)
+        return video.float().clamp_(-1, 1)
+
+    def encode(self, videos, device, tiled=False,end_=False, tile_size=(34, 34), tile_stride=(18, 16)):
+
+        videos = [video.to("cpu") for video in videos]
+        hidden_states = []
+        for video in videos:
+            video = video.unsqueeze(0)
+            if tiled:
+                tile_size = (tile_size[0], tile_size[1])
+                tile_stride = (tile_stride[0], tile_stride[1])
+                hidden_state = self.tiled_encode(video, device, tile_size, tile_stride)
+            else:
+                if end_:
+                    hidden_state = self.double_encode(video, device)
+                else:
+                    hidden_state = self.single_encode(video, device)
+            hidden_state = hidden_state.squeeze(0)
+            hidden_states.append(hidden_state)
+        hidden_states = torch.stack(hidden_states)
+        return hidden_states
+
+
+    def decode(self, hidden_states, device, tiled=False, end_=False, tile_size=(34, 34), tile_stride=(18, 16)):
+        hidden_states = [hidden_state.to("cpu") for hidden_state in hidden_states]
+        videos = []
+        for hidden_state in hidden_states:
+            hidden_state = hidden_state.unsqueeze(0)
+            if tiled:
+                video = self.tiled_decode(hidden_state, device, tile_size, tile_stride)
+            else:
+                if end_:
+                    video = self.double_decode(hidden_state, device)
+                else:
+                    video = self.single_decode(hidden_state, device)
+            video = video.squeeze(0)
+            videos.append(video)
+        return videos
+
+
+    @staticmethod
+    def state_dict_converter():
+        return WanVideoVAEStateDictConverter()
+
+
+class WanVideoVAEStateDictConverter:
+
+    def __init__(self):
+        pass
+
+    def from_civitai(self, state_dict):
+        state_dict_ = {}
+        if 'model_state' in state_dict:
+            state_dict = state_dict['model_state']
+        for name in state_dict:
+            state_dict_['model.' + name] = state_dict[name]
+        return state_dict_