Improvement in WAN VAE Performance - #464
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the Wan pipeline and autoencoder models to support conditional sharding constraints, optimize classifier-free guidance by avoiding redundant latent concatenations, and adapt the caching logic in inference pipelines to be fully compatible with JAX tracing. The review feedback highlights several critical areas for improvement: removing @nnx.jit from VAE encoding/decoding may cause silent performance regressions in the VACE pipeline; the new nearest-neighbor upsampling using jnp.repeat needs to guard against non-integer scale factors; the use of addressable_data should be replaced with standard JAX APIs like addressable_shards to avoid slow fallbacks; and cache_count must be updated within the JAX-compatible caching loops to ensure accurate cache ratio logging.
Perseus14
left a comment
Perseus14
left a comment
There was a problem hiding this comment.
Few minor changes requested.
Please squash the commits.
- Implement spatial/temporal parallelism for WAN VAE - Replace resize with jnp.repeat for upsampling - Resolve eager execution regression in VACE pipeline using vae_encode_pass - Add comprehensive sharding validation and fallbacks in VAE with max_logging - Fix multi-host addressable data logic and TeaCache tracking bugs - Optimize pipeline formatting to reduce host overhead - Restore profiler trace dumping functionality in generate_wan.py
4 participants
Improvement in WAN VAE Performance
Changes made:
Performance Impact:
Tested on v6e-8 for 720p, 81-frame video. These changes reduced step timings from:
Before:
After:
Leading to a 75% improvement in latency.
For v7x-8, the observations are below. The improvement is not as significant as that in v6e-8.
Before:
VAE Decode: 3.6s
- TPU Compute: 2.9s
- Host Formatting: 0.7s
After:
VAE Decode: 2.9s
- TPU Compute: 2.9s
- Host Formatting: 0.0s