Enable JIT Compilation of WAN VAE Encoder/Decoder Forward Passes

[prishajain1]

Refactor WAN VAE to enable JIT compilation of individual encoder and decoder forward passes by converting mutable cache management to immutable functional patterns. Each forward pass through the encoder/decoder network can now be JIT compiled while the overall encode/decode pipeline (chunking, looping) remains outside JIT.

Files modified:

• autoencoder_kl_wan.py :
  • Refactored cache management from mutable lists to immutable tuples for JAX compatibility
  • Added PyTree registration for AutoencoderKLWanCache and RepSentinel to enable passing through JIT boundaries
  • Updated internal module signatures to return tuples (output, feat_cache, feat_idx) for explicit data flow
  • Added @nnx.jit decorators to WanEncoder3d and WanDecoder3d forward passes
  • Optimized cache initialization by moving structure calculations to init()
• vae_flax.py :
  • Added WanDiagonalGaussianDistribution subclass with PyTree registration for JIT-safe distribution handling
  • Inherits all functionality from FlaxDiagonalGaussianDistribution with added JAX transformation support
• base_wan_i2v_14b.yml : configs modified
• base_wan_i2v_27b.yml : configs modified

Tested with:

• WAN 2.1 T2V inference and training
• WAN 2.2 T2V inference
• WAN 2.1 I2V inference
• WAN 2.2 I2V inference

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid