[Feature] Implement tiled VAE encoding/decoding for Wan model. #11414

c8ef · 2025-04-24T14:30:11Z

What does this PR do?

Implement tiled VAE encoding/decoding for Wan model.

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline?
Did you read our philosophy doc (important for complex PRs)?
Was this discussed/approved via a GitHub issue or the forum? 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, and
here are tips on formatting docstrings.
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.

c8ef · 2025-04-24T14:33:50Z

Hi @a-r-r-o-w @yiyixuxu, could you please help reviewing this patch? Thanks!

c8ef

I hope the inline comment makes sense to you~

c8ef · 2025-04-24T14:59:34Z

src/diffusers/models/autoencoders/autoencoder_kl_wan.py

@@ -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] = [


These function parameters are not being used. They have been removed in this patch but can be added back at any time if needed.

c8ef · 2025-04-24T15:00:48Z

tests/models/autoencoders/test_models_autoencoder_wan.py

+        batch_size = 2
+        num_frames = 9
+        num_channels = 3
+        sizes = (640, 480)


Add another input because the (16, 16) tensor is too small for tiling operations.

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.

That's a good idea! Done.

c8ef · 2025-04-24T15:03:34Z

tests/models/autoencoders/test_models_autoencoder_wan.py

+
+        self.assertLess(
+            (output_without_tiling.detach().cpu().numpy() - output_with_tiling.detach().cpu().numpy()).max(),
+            0.5,


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.

a-r-r-o-w

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 🤗

a-r-r-o-w · 2025-04-25T15:21:18Z

src/diffusers/models/autoencoders/autoencoder_kl_wan.py

@@ -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] = [


Indeed it might not be used in the diffusers codebase at the moment, but it is being used downstream in a few repositories (example). 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

I add the parameter back~

a-r-r-o-w · 2025-04-25T15:22:04Z

src/diffusers/models/autoencoders/autoencoder_kl_wan.py

-            2.8251,
-            1.9160,
-        ],
+        spatial_compression_ratio: int = 8,


Let's not add new parameter now because it will lead to unnecessary config warning

a-r-r-o-w · 2025-04-25T15:23:15Z

src/diffusers/models/autoencoders/autoencoder_kl_wan.py

@@ -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


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)

a-r-r-o-w · 2025-04-25T15:23:26Z

src/diffusers/models/autoencoders/autoencoder_kl_wan.py

+        # 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


Let's also add use_slicing

a-r-r-o-w · 2025-04-25T15:25:35Z

tests/models/autoencoders/test_models_autoencoder_wan.py

+        batch_size = 2
+        num_frames = 9
+        num_channels = 3
+        sizes = (640, 480)


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.

a-r-r-o-w · 2025-04-25T15:26:03Z

src/diffusers/models/autoencoders/autoencoder_kl_wan.py

@@ -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


Let's make sure to support use_slicing here as well

Done. Test case added.

c8ef commented Apr 24, 2025

View reviewed changes

yiyixuxu requested a review from a-r-r-o-w April 24, 2025 18:40

a-r-r-o-w reviewed Apr 25, 2025

View reviewed changes

implement tiled encode/decode

822afe0

c8ef force-pushed the main branch from 5878e2f to 822afe0 Compare April 25, 2025 17:02

c8ef requested a review from a-r-r-o-w April 25, 2025 17:04

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Feature] Implement tiled VAE encoding/decoding for Wan model. #11414

[Feature] Implement tiled VAE encoding/decoding for Wan model. #11414

c8ef commented Apr 24, 2025 •

edited

Loading

c8ef commented Apr 24, 2025

c8ef left a comment

c8ef Apr 24, 2025

c8ef Apr 24, 2025

a-r-r-o-w Apr 25, 2025

c8ef Apr 25, 2025

c8ef Apr 24, 2025

a-r-r-o-w left a comment

a-r-r-o-w Apr 25, 2025

c8ef Apr 25, 2025

a-r-r-o-w Apr 25, 2025

c8ef Apr 25, 2025

a-r-r-o-w Apr 25, 2025

c8ef Apr 25, 2025

a-r-r-o-w Apr 25, 2025

c8ef Apr 25, 2025

a-r-r-o-w Apr 25, 2025

a-r-r-o-w Apr 25, 2025

c8ef Apr 25, 2025

[Feature] Implement tiled VAE encoding/decoding for Wan model. #11414

Are you sure you want to change the base?

[Feature] Implement tiled VAE encoding/decoding for Wan model. #11414

Conversation

c8ef commented Apr 24, 2025 • edited Loading

What does this PR do?

Before submitting

Who can review?

c8ef commented Apr 24, 2025

c8ef left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

a-r-r-o-w left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

c8ef commented Apr 24, 2025 •

edited

Loading