[Bug / Optimization] Make AutoencoderKLWan stateless to fix torch.compile compatibility and enable 1.45x speedup

[NotNANtoN]

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

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!