diffusers icon indicating copy to clipboard operation
diffusers copied to clipboard

[WIP] Add parallelized tiled VAE support to AutoencoderKL and AutoencoderKLWan

Open TmacAaron opened this issue 2 months ago • 0 comments

What does this PR do?

When memory is constrained, VAE tiling serves as a critical optimization to reduce memory footprint—here’s how it works at its core:

  1. Split the VAE’s original input (e.g., high-resolution images or latent tensors) into multiple smaller, independently processable tiles;
  2. Feed each tile into the VAE separately for encoding or decoding to generate tile-specific results;
  3. Stitch all processed tiles back to their original positional layout, with overlapping regions optimized via fusion algorithms to ensure the final output is seamless and natural.

However, the default VAE tiling workflow executes tile-wise operations sequentially - even when multiple devices are available. This leads to underutilization of hardware resources, as each tile’s VAE computation is inherently independent and can be parallelized.

To address this inefficiency, this PR introduces parallelized tiled VAE processing for AutoencoderKL and AutoencoderKLWan. By leveraging multi-device capabilities, we distribute the independent tile computations across available hardware, enabling simultaneous processing of multiple tiles. This not only maximizes resource utilization but also accelerates end-to-end VAE encoding/decoding throughput - all while retaining the memory-saving benefits of tiling for memory-constrained environments.

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.

TmacAaron avatar Nov 04 '25 09:11 TmacAaron