Repository navigation
Conversation
The WAN VAE always sharded the width of its activations over the vae_spatial mesh axis. This works for the default 1280x720 landscape size, but users may also want to run inference on portrait videos 720x1280. The VAE downsamples their width to 720, 360, 180 and 90, which isn't always divisible by vae_spatial (e.g., 90 for 4). This change adds a _spatial_sharding(mesh, shape) helper that shards height if width **is not divisible** by vae_spatial, but height is divisible. Videos with a divisible width keep the current behavior.
![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 introduces a helper function _spatial_sharding to dynamically determine the sharding of a 5D array over the vae_spatial axis, replacing hardcoded sharding logic in __call__, _encode, and _decode. The review feedback recommends defensively checking that the input shape is 5D in _spatial_sharding to prevent unpacking errors, and addresses a potential NameError in _decode when spatial_sharding is None.
| if self.mesh is not None and "vae_spatial" in self.mesh.shape | ||
| else None | ||
| ) | ||
| spatial_sharding = _spatial_sharding(self.mesh, x.shape) |
There was a problem hiding this comment.
There is a critical bug in _decode where if spatial_sharding is None (e.g., when running on a single device without a mesh), chunk_in_0 (line 1298) and chunk_in_1 (line 1307) are never defined. This will raise a NameError when they are passed to self.decoder on lines 1299 and 1308.
To fix this, you should define chunk_in_0 and chunk_in_1 unconditionally and apply the sharding constraint conditionally. For example:
# First chunk (i=0)
with jax.named_scope("AutoencoderKLWan_decode_chunk_0"):
chunk_in_0 = x[:, 0:1, ...]
if spatial_sharding is not None:
chunk_in_0 = jax.lax.with_sharding_constraint(chunk_in_0, spatial_sharding)
out_0, dec_feat_map, _ = self.decoder(chunk_in_0, feat_cache=dec_feat_map, feat_idx=0)And similarly for chunk_in_1.
| if mesh is None or "vae_spatial" not in mesh.shape: | ||
| return None | ||
| _, _, height, width, _ = shape |
There was a problem hiding this comment.
To prevent potential ValueError during unpacking, we should defensively check that the input shape is indeed 5D (i.e., has a length of 5) before attempting to unpack it. If it is not 5D, we can safely return None to skip applying the sharding constraint.
| if mesh is None or "vae_spatial" not in mesh.shape: | |
| return None | |
| _, _, height, width, _ = shape | |
| if mesh is None or "vae_spatial" not in mesh.shape or len(shape) != 5: | |
| return None | |
| _, _, height, width, _ = shape |
The WAN VAE always sharded the width of its activations over the vae_spatial mesh axis. This works for the default 1280x720 landscape size, but users may also want to run inference on portrait videos 720x1280. The VAE downsamples their width to 720, 360, 180 and 90, which isn't always divisible by vae_spatial (e.g., 90 for 4).
This change adds a _spatial_sharding(mesh, shape) helper that shards height if width is not divisible by vae_spatial, but height is divisible. Videos with a divisible width keep the current behavior.