Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

c8ef
Copy link

@c8ef c8ef commented Apr 24, 2025

What does this PR do?

Implement tiled VAE encoding/decoding for Wan model.

Before submitting

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
Copy link
Author

c8ef commented Apr 24, 2025

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

Copy link
Author

@c8ef c8ef left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I hope the inline comment makes sense to you~

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

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

batch_size = 2
num_frames = 9
num_channels = 3
sizes = (640, 480)
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


self.assertLess(
(output_without_tiling.detach().cpu().numpy() - output_with_tiling.detach().cpu().numpy()).max(),
0.5,
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@yiyixuxu yiyixuxu requested a review from a-r-r-o-w April 24, 2025 18:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant