huggingface / huggingface/diffusers
[Bug / Optimization] Make AutoencoderKLWan stateless to fix torch.compile compatibility and enable 1.45x speedup
- Dominant language
- Python
- Stars
- 34.5k
- Forks
- 7.3k
- Avg merge
- 3d 3h
- Merged PRs (30d)
- 91
Description
### Description
`AutoencoderKLWan` (used in Wan2.1 and Wan2.2 models) mutates internal instance attributes during its forward pass in `_encode`, `_decode`, `tiled_encode`, and `tiled_decode`:
- `self._enc_feat_map`
- `self._enc_conv_idx`
- `self._feat_map`
- `self._conv_idx`
Furthermore, the intermediate blocks (`WanResample`, `WanResidualBlock`, `WanMidBlock`, `WanResidualDownBlock`, `WanEncoder3d`, `WanResidualUpBlock`, `WanUpBlock`, `WanDecoder3d`) use a mutable default argument `feat_idx=[0]`.
Because module attributes are repeatedly mutated during the forward pass and `self.clear_cache()` reallocates lists on `self`, compiling `AutoencoderKLWan` with `torch.compile` causes graph breaks and prevents efficient compilation.
### Reproduction
```python
import torch
from diffusers import AutoencoderKLWan
vae = AutoencoderKLWan.from_pretrained(
"Wan-AI/Wan2.1-T2V-1.3B", subfolder="vae", torch_dtype=torch.bfloat16
).cuda()
vae = torch.compile(vae)
# Input shape: (B, C, T, H, W) e.g. 5 frames
x = torch.randn(1, 3, 5, 480, 640, dtype=torch.bfloat16, device="cuda")
with torch.no_grad():
# Mutates self._enc_feat_map and self._enc_conv_idx during forward pass
latents = vae.encode(x).latent_dist.sample()
```
### Proposed Solution
Make `AutoencoderKLWan` stateless during execution:
1. Allocate `feat_map = [None] * self._cached_conv_counts["encoder"]` (and `decoder`) locally inside `_encode`, `_decode`, `tiled_encode`, and `tiled_decode`, passing the local list and a local `conv_idx = [0]` to child modules.
2. Change mutable default arguments from `feat_idx=[0]` to `feat_idx=None` (initializing to `[0]` inside the method when `None`).
3. Preserve `clear_cache()` as a method for backward compatibility.
### Verification & Performance
Tested on NVIDIA RTX 4090 (PyTorch 2.11 / CUDA 12.8, BF16, 5 frames 480x640):
- **Numerical Parity**: `max_abs_diff = 0.0` (exact match to uncompiled eager baseline).
- **Unit Tests**: All 32 unit tests in `tests/models/autoencoders/test_models_autoencoder_wan.py` pass cleanly.
- **Speedup**:
- Eager uncompiled: `84.43 ms`
- `torch.compile(mode="default")`: `61.10 ms` (**1.38x speedup**)
- `torch.compile(mode="max-autotune-no-cudagraphs")`: `58.38 ms` (**1.45x speedup**)
We have a clean PR ready with these changes if the maintainers would like us to open it!
Contributor guide
Research direction
Start with AutoencoderKLWan's _encode, _decode, tiled_encode, and tiled_decode methods, then inspect the listed Wan intermediate blocks and tests/models/autoencoders/test_models_autoencoder_wan.py. Run the 32 existing tests first; done means forward execution no longer mutates module state, the tests retain numerical parity, and torch.compile compatibility is verified.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100