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.
| if self.method == "nearest": | ||
| scale_h = int(self.scale_factor[0]) | ||
| scale_w = int(self.scale_factor[1]) | ||
| out = jnp.repeat(jnp.repeat(x, scale_h, axis=1), scale_w, axis=2) | ||
| else: | ||
| out = jax.image.resize(x.astype(jnp.float32), (n, target_h, target_w, c), method=self.method) | ||
| out = out.astype(input_dtype) | ||
| return out |
There was a problem hiding this comment.
If self.scale_factor contains non-integer values (e.g., fractional scaling), using int(self.scale_factor[0]) with jnp.repeat will truncate the scale factor and produce an incorrect output shape, leading to shape mismatches. We should ensure that jnp.repeat is only used when the scale factors are indeed integers, falling back to jax.image.resize otherwise.
| if self.method == "nearest": | |
| scale_h = int(self.scale_factor[0]) | |
| scale_w = int(self.scale_factor[1]) | |
| out = jnp.repeat(jnp.repeat(x, scale_h, axis=1), scale_w, axis=2) | |
| else: | |
| out = jax.image.resize(x.astype(jnp.float32), (n, target_h, target_w, c), method=self.method) | |
| out = out.astype(input_dtype) | |
| return out | |
| if self.method == "nearest" and self.scale_factor[0] == int(self.scale_factor[0]) and self.scale_factor[1] == int(self.scale_factor[1]): | |
| scale_h = int(self.scale_factor[0]) | |
| scale_w = int(self.scale_factor[1]) | |
| out = jnp.repeat(jnp.repeat(x, scale_h, axis=1), scale_w, axis=2) | |
| else: | |
| out = jax.image.resize(x.astype(jnp.float32), (n, target_h, target_w, c), method=self.method) | |
| out = out.astype(input_dtype) | |
| return out |
| if hasattr(video, "addressable_data"): | ||
| video = np.asarray(video.addressable_data(0)) | ||
| else: | ||
| video = np.asarray(video) |
There was a problem hiding this comment.
In standard public JAX, jax.Array does not have an addressable_data method. This will cause hasattr(video, "addressable_data") to always evaluate to False, silently falling back to np.asarray(video) which can trigger slow cross-device transfers on multi-host setups. Consider using addressable_shards[0].data instead, which is the standard JAX API to retrieve the local numpy array of the first addressable shard.
| if hasattr(video, "addressable_data"): | |
| video = np.asarray(video.addressable_data(0)) | |
| else: | |
| video = np.asarray(video) | |
| if hasattr(video, "addressable_shards") and len(video.addressable_shards) > 0: | |
| video = np.asarray(video.addressable_shards[0].data) | |
| else: | |
| video = np.asarray(video) |
…g, teacache stats, and multi-host addressable data
Perseus14
left a comment
Perseus14
left a comment
There was a problem hiding this comment.
Few minor changes requested.
Please squash the commits.
| if axis_name in mesh.shape: | ||
| mesh_axis_size *= mesh.shape[axis_name] | ||
| if x.shape[axis_idx] % mesh_axis_size != 0: | ||
| import logging |
There was a problem hiding this comment.
Let's use our custom max_logging here.
| out = jnp.repeat(jnp.repeat(x, scale_h, axis=1), scale_w, axis=2) | ||
| else: | ||
| if self.method == "nearest": | ||
| import logging |
| ) | ||
| self.mesh = mesh | ||
|
|
||
| @nnx.jit |
There was a problem hiding this comment.
We might need to keep this as WAN VACE directly uses _encode
| return (posterior,) | ||
| return FlaxAutoencoderKLOutput(latent_dist=posterior) | ||
|
|
||
| @nnx.jit |
3 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.