Repository: huggingface/diffusers
Pull request #11414: [Feature] Implement tiled VAE encoding/decoding for Wan model.
Author: login: c8ef
Created: 2025-04-24T14:30:11Z | Merged: 2025-05-05T10:37:15Z

# What does this PR do?

Implement tiled VAE encoding/decoding for Wan model.
<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->



## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/diffusers/blob/main/CONTRIBUTING.md)?
- [x] Did you read our [philosophy doc](https://github.com/huggingface/diffusers/blob/main/PHILOSOPHY.md) (important for complex PRs)?
- [ ] Was this discussed/approved via a GitHub issue or the [forum](https://discuss.huggingface.co/c/discussion-related-to-httpsgithubcomhuggingfacediffusers/63)? Please add a link to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/diffusers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/diffusers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @.

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Core library:

- Schedulers: @yiyixuxu
- Pipelines and pipeline callbacks: @yiyixuxu and @asomoza
- Training examples: @sayakpaul
- Docs: @stevhliu and @sayakpaul
- JAX and MPS: @pcuenca
- Audio: @sanchit-gandhi
- General functionalities: @sayakpaul @yiyixuxu @DN6

Integrations:

- deepspeed: HF Trainer/Accelerate: @SunMarc
- PEFT: @sayakpaul @BenjaminBossan

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- transformers: [different repo](https://github.com/huggingface/transformers)
- safetensors: [different repo](https://github.com/huggingface/safetensors)

-->



Diff:
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py
index fafb1fe867e3..fe00d8c078ff 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py
@@ -730,6 +730,76 @@ def __init__(
             base_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout
         )
 
+        self.spatial_compression_ratio = 2 ** len(self.temperal_downsample)
+
+        # When decoding a batch of video latents at a time, one can save memory by slicing across the batch dimension
+        # to perform decoding of a single video latent at a time.
+        self.use_slicing = False
+
+        # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent
+        # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the
+        # intermediate tiles together, the memory requirement can be lowered.
+        self.use_tiling = False
+
+        # The minimal tile height and width for spatial tiling to be used
+        self.tile_sample_min_height = 256
+        self.tile_sample_min_width = 256
+
+        # The minimal distance between two spatial tiles
+        self.tile_sample_stride_height = 192
+        self.tile_sample_stride_width = 192
+
+    def enable_tiling(
+        self,
+        tile_sample_min_height: Optional[int] = None,
+        tile_sample_min_width: Optional[int] = None,
+        tile_sample_stride_height: Optional[float] = None,
+        tile_sample_stride_width: Optional[float] = None,
+    ) -> None:
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+
+        Args:
+            tile_sample_min_height (`int`, *optional*):
+                The minimum height required for a sample to be separated into tiles across the height dimension.
+            tile_sample_min_width (`int`, *optional*):
+                The minimum width required for a sample to be separated into tiles across the width dimension.
+            tile_sample_stride_height (`int`, *optional*):
+                The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are
+                no tiling artifacts produced across the height dimension.
+            tile_sample_stride_width (`int`, *optional*):
+                The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling
+                artifacts produced across the width dimension.
+        """
+        self.use_tiling = True
+        self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height
+        self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width
+        self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height
+        self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width
+
+    def disable_tiling(self) -> None:
+        r"""
+        Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.use_tiling = False
+
+    def enable_slicing(self) -> None:
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.use_slicing = True
+
+    def disable_slicing(self) -> None:
+        r"""
+        Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing
+        decoding in one step.
+        """
+        self.use_slicing = False
+
     def clear_cache(self):
         def _count_conv3d(model):
             count = 0
@@ -746,11 +816,14 @@ def _count_conv3d(model):
         self._enc_conv_idx = [0]
         self._enc_feat_map = [None] * self._enc_conv_num
 
-    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+    def _encode(self, x: torch.Tensor):
+        _, _, num_frame, height, width = x.shape
+
+        if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height):
+            return self.tiled_encode(x)
+
         self.clear_cache()
-        ## cache
-        t = x.shape[2]
-        iter_ = 1 + (t - 1) // 4
+        iter_ = 1 + (num_frame - 1) // 4
         for i in range(iter_):
             self._enc_conv_idx = [0]
             if i == 0:
@@ -764,8 +837,6 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor:
                 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)
         self.clear_cache()
         return enc
 
@@ -785,18 +856,28 @@ def encode(
                 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.
         """
-        h = self._encode(x)
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self._encode(x)
         posterior = DiagonalGaussianDistribution(h)
+
         if not return_dict:
             return (posterior,)
         return AutoencoderKLOutput(latent_dist=posterior)
 
-    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
-        self.clear_cache()
+    def _decode(self, z: torch.Tensor, return_dict: bool = True):
+        _, _, num_frame, height, width = z.shape
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+
+        if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height):
+            return self.tiled_decode(z, return_dict=return_dict)
 
-        iter_ = z.shape[2]
+        self.clear_cache()
         x = self.post_quant_conv(z)
-        for i in range(iter_):
+        for i in range(num_frame):
             self._conv_idx = [0]
             if i == 0:
                 out = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)
@@ -826,12 +907,161 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutp
                 If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
                 returned.
         """
-        decoded = self._decode(z).sample
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z).sample
+
         if not return_dict:
             return (decoded,)
-
         return DecoderOutput(sample=decoded)
 
+    def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (
+                y / blend_extent
+            )
+        return b
+
+    def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (
+                x / blend_extent
+            )
+        return b
+
+    def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
+        r"""Encode a batch of images using a tiled encoder.
+
+        Args:
+            x (`torch.Tensor`): Input batch of videos.
+
+        Returns:
+            `torch.Tensor`:
+                The latent representation of the encoded videos.
+        """
+        _, _, num_frames, height, width = x.shape
+        latent_height = height // self.spatial_compression_ratio
+        latent_width = width // self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = tile_latent_min_height - tile_latent_stride_height
+        blend_width = tile_latent_min_width - tile_latent_stride_width
+
+        # Split x into overlapping tiles and encode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, self.tile_sample_stride_height):
+            row = []
+            for j in range(0, width, self.tile_sample_stride_width):
+                self.clear_cache()
+                time = []
+                frame_range = 1 + (num_frames - 1) // 4
+                for k in range(frame_range):
+                    self._enc_conv_idx = [0]
+                    if k == 0:
+                        tile = x[:, :, :1, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width]
+                    else:
+                        tile = x[
+                            :,
+                            :,
+                            1 + 4 * (k - 1) : 1 + 4 * k,
+                            i : i + self.tile_sample_min_height,
+                            j : j + self.tile_sample_min_width,
+                        ]
+                    tile = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+                    tile = self.quant_conv(tile)
+                    time.append(tile)
+                row.append(torch.cat(time, dim=2))
+            rows.append(row)
+        self.clear_cache()
+
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_height)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_width)
+                result_row.append(tile[:, :, :, :tile_latent_stride_height, :tile_latent_stride_width])
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
+        return enc
+
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+        r"""
+        Decode a batch of images using a tiled decoder.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not 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
+                returned.
+        """
+        _, _, num_frames, height, width = z.shape
+        sample_height = height * self.spatial_compression_ratio
+        sample_width = width * self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+        blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+        # Split z into overlapping tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, tile_latent_stride_height):
+            row = []
+            for j in range(0, width, tile_latent_stride_width):
+                self.clear_cache()
+                time = []
+                for k in range(num_frames):
+                    self._conv_idx = [0]
+                    tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
+                    tile = self.post_quant_conv(tile)
+                    decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx)
+                    time.append(decoded)
+                row.append(torch.cat(time, dim=2))
+            rows.append(row)
+        self.clear_cache()
+
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_height)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_width)
+                result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width])
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width]
+
+        if not return_dict:
+            return (dec,)
+        return DecoderOutput(sample=dec)
+
     def forward(
         self,
         sample: torch.Tensor,
diff --git a/tests/models/autoencoders/test_models_autoencoder_wan.py b/tests/models/autoencoders/test_models_autoencoder_wan.py
index ffc474039889..777fc56c674b 100644
--- a/tests/models/autoencoders/test_models_autoencoder_wan.py
+++ b/tests/models/autoencoders/test_models_autoencoder_wan.py
@@ -15,6 +15,8 @@
 
 import unittest
 
+import torch
+
 from diffusers import AutoencoderKLWan
 from diffusers.utils.testing_utils import enable_full_determinism, floats_tensor, torch_device
 
@@ -44,9 +46,16 @@ def dummy_input(self):
         num_frames = 9
         num_channels = 3
         sizes = (16, 16)
-
         image = floats_tensor((batch_size, num_channels, num_frames) + sizes).to(torch_device)
+        return {"sample": image}
 
+    @property
+    def dummy_input_tiling(self):
+        batch_size = 2
+        num_frames = 9
+        num_channels = 3
+        sizes = (128, 128)
+        image = floats_tensor((batch_size, num_channels, num_frames) + sizes).to(torch_device)
         return {"sample": image}
 
     @property
@@ -62,6 +71,73 @@ def prepare_init_args_and_inputs_for_common(self):
         inputs_dict = self.dummy_input
         return init_dict, inputs_dict
 
+    def prepare_init_args_and_inputs_for_tiling(self):
+        init_dict = self.get_autoencoder_kl_wan_config()
+        inputs_dict = self.dummy_input_tiling
+        return init_dict, inputs_dict
+
+    def test_enable_disable_tiling(self):
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_tiling()
+
+        torch.manual_seed(0)
+        model = self.model_class(**init_dict).to(torch_device)
+
+        inputs_dict.update({"return_dict": False})
+
+        torch.manual_seed(0)
+        output_without_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        torch.manual_seed(0)
+        model.enable_tiling(96, 96, 64, 64)
+        output_with_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        self.assertLess(
+            (output_without_tiling.detach().cpu().numpy() - output_with_tiling.detach().cpu().numpy()).max(),
+            0.5,
+            "VAE tiling should not affect the inference results",
+        )
+
+        torch.manual_seed(0)
+        model.disable_tiling()
+        output_without_tiling_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        self.assertEqual(
+            output_without_tiling.detach().cpu().numpy().all(),
+            output_without_tiling_2.detach().cpu().numpy().all(),
+            "Without tiling outputs should match with the outputs when tiling is manually disabled.",
+        )
+
+    def test_enable_disable_slicing(self):
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+
+        torch.manual_seed(0)
+        model = self.model_class(**init_dict).to(torch_device)
+
+        inputs_dict.update({"return_dict": False})
+
+        torch.manual_seed(0)
+        output_without_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        torch.manual_seed(0)
+        model.enable_slicing()
+        output_with_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        self.assertLess(
+            (output_without_slicing.detach().cpu().numpy() - output_with_slicing.detach().cpu().numpy()).max(),
+            0.05,
+            "VAE slicing should not affect the inference results",
+        )
+
+        torch.manual_seed(0)
+        model.disable_slicing()
+        output_without_slicing_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        self.assertEqual(
+            output_without_slicing.detach().cpu().numpy().all(),
+            output_without_slicing_2.detach().cpu().numpy().all(),
+            "Without slicing outputs should match with the outputs when slicing is manually disabled.",
+        )
+
     @unittest.skip("Gradient checkpointing has not been implemented yet")
     def test_gradient_checkpointing_is_applied(self):
         pass


Reviews:
user: login: c8ef
body: I hope the inline comment makes sense to you~
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T15:04:13Z
user: login: a-r-r-o-w
body: Thanks for working on this @c8ef! Most of the changes look good to me. Will have to verify the visual quality once before we merge -- I can look into that after some of the reviews here are addressed 🤗 
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T15:28:05Z
user: login: c8ef
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T17:02:50Z
user: login: c8ef
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T17:03:06Z
user: login: c8ef
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T17:03:13Z
user: login: c8ef
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T17:03:19Z
user: login: c8ef
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T17:03:27Z
user: login: c8ef
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T17:03:38Z
user: login: a-r-r-o-w
body: Thanks for the updates after last review! Most changes lgtm but some further refactoring is needed
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T01:23:39Z
user: login: c8ef
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T02:04:13Z
user: login: c8ef
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T14:33:50Z
user: login: c8ef
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T14:34:40Z
user: login: a-r-r-o-w
body: Thanks for addressing the reviews, really awesome work!

Visually, the code looks correct to me, but I need to verify the results before/after numerically to know for sure that there is no regression in the behaviour. Currently I'm super busy with something but I'll try to merge this after the weekend (sorry for the delay!)
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-02T03:19:17Z
user: login: a-r-r-o-w
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-05T09:06:38Z
user: login: a-r-r-o-w
body: Thanks for the awesome work! The changes are good to merge

Here's a minimal script I used for testing:

<details>
<summary> code </summary>

```python
import argparse

import torch
from diffusers import AutoencoderKLWan
from diffusers.video_processor import VideoProcessor
from diffusers.utils import export_to_video, load_video

@torch.no_grad()
def main(args):
    height = 480
    width = 768

    torch.cuda.reset_peak_memory_stats()

    vae = AutoencoderKLWan.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers", subfolder="vae", torch_dtype=torch.bfloat16)
    vae.to("cuda")
    if args.enable_slicing:
        vae.enable_slicing()
    if args.enable_tiling:
        vae.enable_tiling()
    video_processor = VideoProcessor(vae_latent_channels=vae.config.z_dim)
    
    video = load_video("inputs/peter-dance.mp4")[::2][:81]
    batch_size = 1
    if args.create_video_batch:
        batch_size = 2
        video = [video] * batch_size
    video = video_processor.preprocess_video(video, height, width)
    video = video.to("cuda", dtype=torch.bfloat16)

    encoded = vae.encode(video).latent_dist.sample(generator=torch.Generator().manual_seed(42))
    print(f"Encoded shape: {encoded.shape}")
    print(f"Max memory (encode): {torch.cuda.max_memory_allocated() / 1024**3:.3f} GB")
    torch.cuda.reset_peak_memory_stats()
    
    decoded = vae.decode(encoded).sample
    print(f"Decoded shape: {decoded.shape}")
    print(f"Max memory (decode): {torch.cuda.max_memory_allocated() / 1024**3:.3f} GB")
    
    videos = video_processor.postprocess_video(decoded, output_type="pil")
    for i in range(batch_size):
        export_to_video(videos[i], f"output{i}.mp4", fps=16)


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--enable_tiling", action="store_true")
    parser.add_argument("--enable_slicing", action="store_true")
    parser.add_argument("--create_video_batch", action="store_true")
    return parser.parse_args()


if __name__ == "__main__":
    args = get_args()
    main(args)
```
</details>

```
(nightly-venv) (nightly-venv) aryan@hf-dgx-01:~/work/diffusers$ python3 dump13.py 
Encoded shape: torch.Size([1, 16, 21, 60, 96])
Max memory (encode): 2.831 GB
Decoded shape: torch.Size([1, 3, 81, 480, 768])
Max memory (decode): 4.451 GB
(nightly-venv) (nightly-venv) aryan@hf-dgx-01:~/work/diffusers$ python3 dump13.py --enable_tiling
Encoded shape: torch.Size([1, 16, 21, 60, 96])
Max memory (encode): 0.846 GB
Decoded shape: torch.Size([1, 3, 81, 480, 768])
Max memory (decode): 1.298 GB
(nightly-venv) (nightly-venv) aryan@hf-dgx-01:~/work/diffusers$ python3 dump13.py --enable_tiling --create_video_batch
Encoded shape: torch.Size([2, 16, 21, 60, 96])
Max memory (encode): 1.453 GB
Decoded shape: torch.Size([2, 3, 81, 480, 768])
Max memory (decode): 2.358 GB
(nightly-venv) (nightly-venv) aryan@hf-dgx-01:~/work/diffusers$ python3 dump13.py --enable_tiling --create_video_batch --enable_slicing
Encoded shape: torch.Size([2, 16, 21, 60, 96])
Max memory (encode): 1.020 GB
Decoded shape: torch.Size([2, 3, 81, 480, 768])
Max memory (decode): 1.637 GB
```

The results are as expected. Additionally, the outputs without any tiling/slicing on this branch match that of main. I'll merge once the tests pass after the latest suggested are applied.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-05T09:54:09Z
user: login: c8ef
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-05T10:04:01Z
user: login: c8ef
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-05T10:04:06Z
user: login: c8ef
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-05T10:04:17Z

Review comments:
diff_hunk: @@ -677,42 +677,7 @@ def __init__(
         attn_scales: List[float] = [],
         temperal_downsample: List[bool] = [False, True, True],
         dropout: float = 0.0,
-        latents_mean: List[float] = [
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: c8ef
body: These function parameters are not being used. They have been removed in this patch but can be added back at any time if needed.
created_at: 2025-04-24T14:59:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -44,9 +46,16 @@ def dummy_input(self):
         num_frames = 9
         num_channels = 3
         sizes = (16, 16)
-
         image = floats_tensor((batch_size, num_channels, num_frames) + sizes).to(torch_device)
+        return {"sample": image}
 
+    @property
+    def dummy_input_tiling(self):
+        batch_size = 2
+        num_frames = 9
+        num_channels = 3
+        sizes = (640, 480)
path: tests/models/autoencoders/test_models_autoencoder_wan.py
user: login: c8ef
body: Add another input because the (16, 16) tensor is too small for tiling operations.
created_at: 2025-04-24T15:00:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -62,6 +71,42 @@ def prepare_init_args_and_inputs_for_common(self):
         inputs_dict = self.dummy_input
         return init_dict, inputs_dict
 
+    def prepare_init_args_and_inputs_for_tiling(self):
+        init_dict = self.get_autoencoder_kl_wan_config()
+        inputs_dict = self.dummy_input_tiling
+        return init_dict, inputs_dict
+
+    def test_enable_disable_tiling(self):
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_tiling()
+
+        torch.manual_seed(0)
+        model = self.model_class(**init_dict).to(torch_device)
+
+        inputs_dict.update({"return_dict": False})
+
+        torch.manual_seed(0)
+        output_without_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        torch.manual_seed(0)
+        model.enable_tiling()
+        output_with_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0]
+
+        self.assertLess(
+            (output_without_tiling.detach().cpu().numpy() - output_with_tiling.detach().cpu().numpy()).max(),
+            0.5,
path: tests/models/autoencoders/test_models_autoencoder_wan.py
user: login: c8ef
body: On my machine, this value is approximately 0.404, and IIRC the average absolute value of these arrays is less than 0.01, which makes me confident that the implementation is correct at some point.
created_at: 2025-04-24T15:03:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -677,42 +677,7 @@ def __init__(
         attn_scales: List[float] = [],
         temperal_downsample: List[bool] = [False, True, True],
         dropout: float = 0.0,
-        latents_mean: List[float] = [
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: a-r-r-o-w
body: Indeed it might not be used in the diffusers codebase at the moment, but it is being used downstream in a few repositories ([example](https://github.com/a-r-r-o-w/finetrainers/blob/c075defc365b5b6ad585251237d83a9976e59620/finetrainers/models/wan/base_specification.py#L78)). Removing this will break downstream so we should keep this anyway.

What you could instead do here to reduce LOC is wrap these two parameters in a non-format block and condense the list into a single line:
```
# fmt: off
latents_mean: List[float] = ...
latents_std: List[float] = ...
# fmt: on
```
created_at: 2025-04-25T15:21:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -677,42 +677,7 @@ def __init__(
         attn_scales: List[float] = [],
         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,
-        ],
-        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,
-        ],
+        spatial_compression_ratio: int = 8,
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: a-r-r-o-w
body: Let's not add new parameter now because it will lead to unnecessary config warning
created_at: 2025-04-25T15:22:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -730,6 +695,58 @@ def __init__(
             base_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout
         )
 
+        self.spatial_compression_ratio = spatial_compression_ratio
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: a-r-r-o-w
body: Instead, let's set this attribute based on the init parameters. The same logic as used in the pipeline can be applied here.

Let's also add temporal_compression_ratio as 2 raised to power of (the number of true values in temporal downsample)
created_at: 2025-04-25T15:23:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -730,6 +695,58 @@ def __init__(
             base_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout
         )
 
+        self.spatial_compression_ratio = spatial_compression_ratio
+
+        # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent
+        # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the
+        # intermediate tiles together, the memory requirement can be lowered.
+        self.use_tiling = False
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: a-r-r-o-w
body: Let's also add use_slicing
created_at: 2025-04-25T15:23:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -44,9 +46,16 @@ def dummy_input(self):
         num_frames = 9
         num_channels = 3
         sizes = (16, 16)
-
         image = floats_tensor((batch_size, num_channels, num_frames) + sizes).to(torch_device)
+        return {"sample": image}
 
+    @property
+    def dummy_input_tiling(self):
+        batch_size = 2
+        num_frames = 9
+        num_channels = 3
+        sizes = (640, 480)
path: tests/models/autoencoders/test_models_autoencoder_wan.py
user: login: a-r-r-o-w
body: Let's try to reduce the size as much as possible because these tests should not cause unexpected slowdowns in the CI. While enabling tiling, you can set different tile width/height and stride, than the default 256 and 192.

(128, 128) would be good, with tile size being `96, 96` and stride being `64, 64`, or something similar/lower.
created_at: 2025-04-25T15:25:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -785,7 +802,11 @@ def encode(
                 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.
         """
-        h = self._encode(x)
+        _, _, _, height, width = x.shape
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: a-r-r-o-w
body: Let's make sure to support use_slicing here as well
created_at: 2025-04-25T15:26:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -44,9 +46,16 @@ def dummy_input(self):
         num_frames = 9
         num_channels = 3
         sizes = (16, 16)
-
         image = floats_tensor((batch_size, num_channels, num_frames) + sizes).to(torch_device)
+        return {"sample": image}
 
+    @property
+    def dummy_input_tiling(self):
+        batch_size = 2
+        num_frames = 9
+        num_channels = 3
+        sizes = (640, 480)
path: tests/models/autoencoders/test_models_autoencoder_wan.py
user: login: c8ef
body: That's a good idea! Done.
created_at: 2025-04-25T17:02:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -677,42 +677,7 @@ def __init__(
         attn_scales: List[float] = [],
         temperal_downsample: List[bool] = [False, True, True],
         dropout: float = 0.0,
-        latents_mean: List[float] = [
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: c8ef
body: I add the parameter back~
created_at: 2025-04-25T17:03:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -677,42 +677,7 @@ def __init__(
         attn_scales: List[float] = [],
         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,
-        ],
-        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,
-        ],
+        spatial_compression_ratio: int = 8,
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: c8ef
body: Done.
created_at: 2025-04-25T17:03:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -730,6 +695,58 @@ def __init__(
             base_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout
         )
 
+        self.spatial_compression_ratio = spatial_compression_ratio
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: c8ef
body: Done.
created_at: 2025-04-25T17:03:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -730,6 +695,58 @@ def __init__(
             base_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout
         )
 
+        self.spatial_compression_ratio = spatial_compression_ratio
+
+        # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent
+        # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the
+        # intermediate tiles together, the memory requirement can be lowered.
+        self.use_tiling = False
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: c8ef
body: Done.
created_at: 2025-04-25T17:03:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -785,7 +802,11 @@ def encode(
                 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.
         """
-        h = self._encode(x)
+        _, _, _, height, width = x.shape
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: c8ef
body: Done. Test case added.
created_at: 2025-04-25T17:03:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -746,7 +817,7 @@ def _count_conv3d(model):
         self._enc_conv_idx = [0]
         self._enc_feat_map = [None] * self._enc_conv_num
 
-    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+    def vanilla_encode(self, x: torch.Tensor) -> torch.Tensor:
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: a-r-r-o-w
body: This logic should not be a separate method. Previously, the behaviour was to perform temporal tiling by default. This was incorrectly done and should've been opt-in but we merged the Wan changes without this being addressed. So, now we need to maintain it that way by having temporal tiling being opt-in by default for both cases (whether `enable_tiling` has actually been called by user or not)

The way I would rewrite this would be:
- Let `encode`/`decode` handle the batch-wise slicing if enabled (your implementation is correct in this regard)
- Let `_encode`/`_decode` handle the spatial tiling based on whether user enabled it or not (your implementation is correct in this regard)
- Instead of the `vanilla_*` methods, move the temporal tiling logic into `_encode`, `_decode`, `tiled_encode`, `tiled_decode`, where temporal tiling should be enabled by default to maintain backwards compatiblity (you can introduce two other flags called `use_framewise_encoding` and `use_framewise_decoding` that are True by default. For when they are False, feel free to simply raise a NotImplementedError unless you'd like to give that a try too :) [ref](https://github.com/huggingface/diffusers/blob/a7e9f85e21dde12f2f2489702ea82db40ebb31d2/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py#L765), [ref2](https://github.com/huggingface/diffusers/blob/a7e9f85e21dde12f2f2489702ea82db40ebb31d2/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py#L1007)


created_at: 2025-04-28T01:22:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -746,7 +817,7 @@ def _count_conv3d(model):
         self._enc_conv_idx = [0]
         self._enc_feat_map = [None] * self._enc_conv_num
 
-    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+    def vanilla_encode(self, x: torch.Tensor) -> torch.Tensor:
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: c8ef
body: Thank you for your detailed guidance! I learned a lot from it and will refactor the code in the coming week.
created_at: 2025-04-28T02:04:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -746,7 +817,7 @@ def _count_conv3d(model):
         self._enc_conv_idx = [0]
         self._enc_feat_map = [None] * self._enc_conv_num
 
-    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+    def vanilla_encode(self, x: torch.Tensor) -> torch.Tensor:
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: c8ef
body: Based on review feedback, I removed the `vanilla_*` function and transferred the logic to `_encode` and `_decode`. Initially, I wanted to create something similar to `temporal_tiled_decode` like Hunyuan VAE, but I discovered a small problem with the feat cache mechanism in Wan VAE when using this approach (or maybe I misunderstood something). As a result, I modified the implementation to resemble cogvideo. The latest patch also eliminates some unused variables and functions that were added in a previous patch.
created_at: 2025-04-28T14:33:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -764,9 +837,6 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor:
                 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)
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: c8ef
body: This is a no-op.
created_at: 2025-04-28T14:34:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -764,9 +837,6 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor:
                 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)
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: a-r-r-o-w
body: Nice find!
created_at: 2025-05-05T09:06:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -764,9 +837,6 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor:
                 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)
-        self.clear_cache()
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: a-r-r-o-w
body: I think this call to `self.clear_cache()` still needs to be present
created_at: 2025-05-05T09:47:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -826,12 +906,159 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutp
                 If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
                 returned.
         """
-        decoded = self._decode(z).sample
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z).sample
+
         if not return_dict:
             return (decoded,)
-
         return DecoderOutput(sample=decoded)
 
+    def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (
+                y / blend_extent
+            )
+        return b
+
+    def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (
+                x / blend_extent
+            )
+        return b
+
+    def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
+        r"""Encode a batch of images using a tiled encoder.
+
+        Args:
+            x (`torch.Tensor`): Input batch of videos.
+
+        Returns:
+            `torch.Tensor`:
+                The latent representation of the encoded videos.
+        """
+        _, _, num_frames, height, width = x.shape
+        latent_height = height // self.spatial_compression_ratio
+        latent_width = width // self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = tile_latent_min_height - tile_latent_stride_height
+        blend_width = tile_latent_min_width - tile_latent_stride_width
+
+        # Split x into overlapping tiles and encode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, self.tile_sample_stride_height):
+            row = []
+            for j in range(0, width, self.tile_sample_stride_width):
+                self.clear_cache()
+                time = []
+                frame_range = 1 + (num_frames - 1) // 4
+                for k in range(frame_range):
+                    self._enc_conv_idx = [0]
+                    if k == 0:
+                        tile = x[:, :, :1, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width]
+                    else:
+                        tile = x[
+                            :,
+                            :,
+                            1 + 4 * (k - 1) : 1 + 4 * k,
+                            i : i + self.tile_sample_min_height,
+                            j : j + self.tile_sample_min_width,
+                        ]
+                    tile = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+                    tile = self.quant_conv(tile)
+                    time.append(tile)
+                row.append(torch.cat(time, dim=2))
+            rows.append(row)
+
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_height)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_width)
+                result_row.append(tile[:, :, :, :tile_latent_stride_height, :tile_latent_stride_width])
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
+        return enc
+
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+        r"""
+        Decode a batch of images using a tiled decoder.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not 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
+                returned.
+        """
+        _, _, num_frames, height, width = z.shape
+        sample_height = height * self.spatial_compression_ratio
+        sample_width = width * self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+        blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+        # Split z into overlapping tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, tile_latent_stride_height):
+            row = []
+            for j in range(0, width, tile_latent_stride_width):
+                self.clear_cache()
+                time = []
+                for k in range(num_frames):
+                    self._conv_idx = [0]
+                    tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
+                    tile = self.post_quant_conv(tile)
+                    decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx)
+                    time.append(decoded)
+                row.append(torch.cat(time, dim=2))
+            rows.append(row)
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: a-r-r-o-w
body: ```suggestion
            rows.append(row)
        self.clear_cache()
```
created_at: 2025-05-05T09:50:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -826,12 +906,159 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutp
                 If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
                 returned.
         """
-        decoded = self._decode(z).sample
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z).sample
+
         if not return_dict:
             return (decoded,)
-
         return DecoderOutput(sample=decoded)
 
+    def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (
+                y / blend_extent
+            )
+        return b
+
+    def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (
+                x / blend_extent
+            )
+        return b
+
+    def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
+        r"""Encode a batch of images using a tiled encoder.
+
+        Args:
+            x (`torch.Tensor`): Input batch of videos.
+
+        Returns:
+            `torch.Tensor`:
+                The latent representation of the encoded videos.
+        """
+        _, _, num_frames, height, width = x.shape
+        latent_height = height // self.spatial_compression_ratio
+        latent_width = width // self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = tile_latent_min_height - tile_latent_stride_height
+        blend_width = tile_latent_min_width - tile_latent_stride_width
+
+        # Split x into overlapping tiles and encode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, self.tile_sample_stride_height):
+            row = []
+            for j in range(0, width, self.tile_sample_stride_width):
+                self.clear_cache()
+                time = []
+                frame_range = 1 + (num_frames - 1) // 4
+                for k in range(frame_range):
+                    self._enc_conv_idx = [0]
+                    if k == 0:
+                        tile = x[:, :, :1, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width]
+                    else:
+                        tile = x[
+                            :,
+                            :,
+                            1 + 4 * (k - 1) : 1 + 4 * k,
+                            i : i + self.tile_sample_min_height,
+                            j : j + self.tile_sample_min_width,
+                        ]
+                    tile = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+                    tile = self.quant_conv(tile)
+                    time.append(tile)
+                row.append(torch.cat(time, dim=2))
+            rows.append(row)
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: a-r-r-o-w
body: ```suggestion
            rows.append(row)
        self.clear_cache()
```
created_at: 2025-05-05T09:51:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -764,9 +837,6 @@ def _encode(self, x: torch.Tensor) -> torch.Tensor:
                 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)
-        self.clear_cache()
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: c8ef
body: Done.
created_at: 2025-05-05T10:04:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -826,12 +906,159 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutp
                 If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
                 returned.
         """
-        decoded = self._decode(z).sample
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z).sample
+
         if not return_dict:
             return (decoded,)
-
         return DecoderOutput(sample=decoded)
 
+    def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (
+                y / blend_extent
+            )
+        return b
+
+    def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (
+                x / blend_extent
+            )
+        return b
+
+    def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
+        r"""Encode a batch of images using a tiled encoder.
+
+        Args:
+            x (`torch.Tensor`): Input batch of videos.
+
+        Returns:
+            `torch.Tensor`:
+                The latent representation of the encoded videos.
+        """
+        _, _, num_frames, height, width = x.shape
+        latent_height = height // self.spatial_compression_ratio
+        latent_width = width // self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = tile_latent_min_height - tile_latent_stride_height
+        blend_width = tile_latent_min_width - tile_latent_stride_width
+
+        # Split x into overlapping tiles and encode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, self.tile_sample_stride_height):
+            row = []
+            for j in range(0, width, self.tile_sample_stride_width):
+                self.clear_cache()
+                time = []
+                frame_range = 1 + (num_frames - 1) // 4
+                for k in range(frame_range):
+                    self._enc_conv_idx = [0]
+                    if k == 0:
+                        tile = x[:, :, :1, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width]
+                    else:
+                        tile = x[
+                            :,
+                            :,
+                            1 + 4 * (k - 1) : 1 + 4 * k,
+                            i : i + self.tile_sample_min_height,
+                            j : j + self.tile_sample_min_width,
+                        ]
+                    tile = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+                    tile = self.quant_conv(tile)
+                    time.append(tile)
+                row.append(torch.cat(time, dim=2))
+            rows.append(row)
+
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_height)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_width)
+                result_row.append(tile[:, :, :, :tile_latent_stride_height, :tile_latent_stride_width])
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
+        return enc
+
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+        r"""
+        Decode a batch of images using a tiled decoder.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not 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
+                returned.
+        """
+        _, _, num_frames, height, width = z.shape
+        sample_height = height * self.spatial_compression_ratio
+        sample_width = width * self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+        blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+        # Split z into overlapping tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, tile_latent_stride_height):
+            row = []
+            for j in range(0, width, tile_latent_stride_width):
+                self.clear_cache()
+                time = []
+                for k in range(num_frames):
+                    self._conv_idx = [0]
+                    tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
+                    tile = self.post_quant_conv(tile)
+                    decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx)
+                    time.append(decoded)
+                row.append(torch.cat(time, dim=2))
+            rows.append(row)
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: c8ef
body: Done.
created_at: 2025-05-05T10:04:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -826,12 +906,159 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutp
                 If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
                 returned.
         """
-        decoded = self._decode(z).sample
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z).sample
+
         if not return_dict:
             return (decoded,)
-
         return DecoderOutput(sample=decoded)
 
+    def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (
+                y / blend_extent
+            )
+        return b
+
+    def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (
+                x / blend_extent
+            )
+        return b
+
+    def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
+        r"""Encode a batch of images using a tiled encoder.
+
+        Args:
+            x (`torch.Tensor`): Input batch of videos.
+
+        Returns:
+            `torch.Tensor`:
+                The latent representation of the encoded videos.
+        """
+        _, _, num_frames, height, width = x.shape
+        latent_height = height // self.spatial_compression_ratio
+        latent_width = width // self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = tile_latent_min_height - tile_latent_stride_height
+        blend_width = tile_latent_min_width - tile_latent_stride_width
+
+        # Split x into overlapping tiles and encode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, self.tile_sample_stride_height):
+            row = []
+            for j in range(0, width, self.tile_sample_stride_width):
+                self.clear_cache()
+                time = []
+                frame_range = 1 + (num_frames - 1) // 4
+                for k in range(frame_range):
+                    self._enc_conv_idx = [0]
+                    if k == 0:
+                        tile = x[:, :, :1, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width]
+                    else:
+                        tile = x[
+                            :,
+                            :,
+                            1 + 4 * (k - 1) : 1 + 4 * k,
+                            i : i + self.tile_sample_min_height,
+                            j : j + self.tile_sample_min_width,
+                        ]
+                    tile = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+                    tile = self.quant_conv(tile)
+                    time.append(tile)
+                row.append(torch.cat(time, dim=2))
+            rows.append(row)
path: src/diffusers/models/autoencoders/autoencoder_kl_wan.py
user: login: c8ef
body: Done.
created_at: 2025-05-05T10:04:17Z
author_association: CONTRIBUTOR

PR comments:
user: login: c8ef
created_at: 2025-04-24T14:33:50Z
body: Hi @a-r-r-o-w @yiyixuxu, could you please help reviewing this patch? Thanks!
author_association: CONTRIBUTOR
user: login: c8ef
created_at: 2025-04-27T14:24:10Z
body: Gentle ping~
author_association: CONTRIBUTOR
user: login: HuggingFaceDocBuilderDev
created_at: 2025-04-28T01:15:51Z
body: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/diffusers/pr_11414). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.
author_association: NONE
user: login: c8ef
created_at: 2025-05-01T09:38:50Z
body: Hi @a-r-r-o-w, could you please take another look?
author_association: CONTRIBUTOR
user: login: c8ef
created_at: 2025-05-05T10:08:49Z
body: > Thanks for the awesome work! The changes are good to merge
> 
> Here's a minimal script I used for testing:
> 
> code
> ```python
> import argparse
> 
> import torch
> from diffusers import AutoencoderKLWan
> from diffusers.video_processor import VideoProcessor
> from diffusers.utils import export_to_video, load_video
> 
> @torch.no_grad()
> def main(args):
>     height = 480
>     width = 768
> 
>     torch.cuda.reset_peak_memory_stats()
> 
>     vae = AutoencoderKLWan.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers", subfolder="vae", torch_dtype=torch.bfloat16)
>     vae.to("cuda")
>     if args.enable_slicing:
>         vae.enable_slicing()
>     if args.enable_tiling:
>         vae.enable_tiling()
>     video_processor = VideoProcessor(vae_latent_channels=vae.config.z_dim)
>     
>     video = load_video("inputs/peter-dance.mp4")[::2][:81]
>     batch_size = 1
>     if args.create_video_batch:
>         batch_size = 2
>         video = [video] * batch_size
>     video = video_processor.preprocess_video(video, height, width)
>     video = video.to("cuda", dtype=torch.bfloat16)
> 
>     encoded = vae.encode(video).latent_dist.sample(generator=torch.Generator().manual_seed(42))
>     print(f"Encoded shape: {encoded.shape}")
>     print(f"Max memory (encode): {torch.cuda.max_memory_allocated() / 1024**3:.3f} GB")
>     torch.cuda.reset_peak_memory_stats()
>     
>     decoded = vae.decode(encoded).sample
>     print(f"Decoded shape: {decoded.shape}")
>     print(f"Max memory (decode): {torch.cuda.max_memory_allocated() / 1024**3:.3f} GB")
>     
>     videos = video_processor.postprocess_video(decoded, output_type="pil")
>     for i in range(batch_size):
>         export_to_video(videos[i], f"output{i}.mp4", fps=16)
> 
> 
> def get_args():
>     parser = argparse.ArgumentParser()
>     parser.add_argument("--enable_tiling", action="store_true")
>     parser.add_argument("--enable_slicing", action="store_true")
>     parser.add_argument("--create_video_batch", action="store_true")
>     return parser.parse_args()
> 
> 
> if __name__ == "__main__":
>     args = get_args()
>     main(args)
> ```
> 
> ```
> (nightly-venv) (nightly-venv) aryan@hf-dgx-01:~/work/diffusers$ python3 dump13.py 
> Encoded shape: torch.Size([1, 16, 21, 60, 96])
> Max memory (encode): 2.831 GB
> Decoded shape: torch.Size([1, 3, 81, 480, 768])
> Max memory (decode): 4.451 GB
> (nightly-venv) (nightly-venv) aryan@hf-dgx-01:~/work/diffusers$ python3 dump13.py --enable_tiling
> Encoded shape: torch.Size([1, 16, 21, 60, 96])
> Max memory (encode): 0.846 GB
> Decoded shape: torch.Size([1, 3, 81, 480, 768])
> Max memory (decode): 1.298 GB
> (nightly-venv) (nightly-venv) aryan@hf-dgx-01:~/work/diffusers$ python3 dump13.py --enable_tiling --create_video_batch
> Encoded shape: torch.Size([2, 16, 21, 60, 96])
> Max memory (encode): 1.453 GB
> Decoded shape: torch.Size([2, 3, 81, 480, 768])
> Max memory (decode): 2.358 GB
> (nightly-venv) (nightly-venv) aryan@hf-dgx-01:~/work/diffusers$ python3 dump13.py --enable_tiling --create_video_batch --enable_slicing
> Encoded shape: torch.Size([2, 16, 21, 60, 96])
> Max memory (encode): 1.020 GB
> Decoded shape: torch.Size([2, 3, 81, 480, 768])
> Max memory (decode): 1.637 GB
> ```
> 
> The results are as expected. Additionally, the outputs without any tiling/slicing on this branch match that of main. I'll merge once the tests pass after the latest suggested are applied.

Thank you for your thorough testing! I appreciate it, as these code snippets are a great starting point for me to delve deeper into the details. Thanks again!
I also believe that tiled encoding/decoding is beneficial because it simplifies distributing the computation across multiple GPUs.
The review comments have been addressed, and it seems that the previous CI failure is unrelated to this patch.
author_association: CONTRIBUTOR