From 85944f7992676f2c1fc027703bed91a17bf6cdc2 Mon Sep 17 00:00:00 2001 From: Nina Shvetsova Date: Wed, 11 Nov 2026 15:48:16 +0000 Subject: [PATCH] Fall back to sharding the height in the WAN VAE. 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. --- .../models/wan/autoencoder_kl_wan.py | 30 +++++++++++-------- 1 file changed, 18 insertions(+), 12 deletions(-) diff --git a/src/maxdiffusion/models/wan/autoencoder_kl_wan.py b/src/maxdiffusion/models/wan/autoencoder_kl_wan.py index 999f5234b..7d3bee300 100644 --- a/src/maxdiffusion/models/wan/autoencoder_kl_wan.py +++ b/src/maxdiffusion/models/wan/autoencoder_kl_wan.py @@ -48,6 +48,21 @@ def _update_cache(cache, idx, value): return cache[:idx] + (value,) + cache[idx + 1 :] +def _spatial_sharding(mesh: Optional[jax.sharding.Mesh], shape: Sequence[int]) -> Optional[NamedSharding]: + """Returns the sharding of a (N, D, H, W, C) array over the vae_spatial axis. + + Shards the width, or the height if only the height is divisible by + vae_spatial. If neither is, the width is still sharded and XLA pads it. + """ + if mesh is None or "vae_spatial" not in mesh.shape: + return None + _, _, height, width, _ = shape + num_shards = mesh.shape["vae_spatial"] + if width % num_shards != 0 and height % num_shards == 0: + return NamedSharding(mesh, P("redundant", None, "vae_spatial", None, None)) + return NamedSharding(mesh, P("redundant", None, None, "vae_spatial", None)) + + # Helper to ensure kernel_size, stride, padding are tuples of 3 integers def _canonicalize_tuple(x: Union[int, Sequence[int]], rank: int, name: str) -> Tuple[int, ...]: """Canonicalizes a value to a tuple of integers.""" @@ -121,7 +136,7 @@ def __init__( def __call__(self, x: jax.Array, cache_x: Optional[jax.Array] = None, idx=-1) -> jax.Array: if self.mesh is not None and "vae_spatial" in self.mesh.shape: - spatial_sharding = NamedSharding(self.mesh, P("redundant", None, None, "vae_spatial", None)) + spatial_sharding = _spatial_sharding(self.mesh, x.shape) if spatial_sharding is not None: x = jax.lax.with_sharding_constraint(x, spatial_sharding) @@ -1169,12 +1184,6 @@ def _encode(self, x: jax.Array, feat_cache: AutoencoderKLWanCache): iter_ = 1 + ((t - 1 + CHUNK_SIZE - 1) // CHUNK_SIZE) if t > 1 else 1 enc_feat_map = feat_cache._enc_feat_map - spatial_sharding = ( - NamedSharding(self.mesh, P("redundant", None, None, "vae_spatial", None)) - if self.mesh is not None and "vae_spatial" in self.mesh.shape - else None - ) - def finalize(out, enc_feat_map): feat_cache._enc_feat_map = enc_feat_map enc = self.quant_conv(out) @@ -1184,6 +1193,7 @@ def finalize(out, enc_feat_map): with jax.named_scope("AutoencoderKLWan_encode_chunk_0"): chunk_0 = x[:, :1, ...] out_0, enc_feat_map, _ = self.encoder(chunk_0, feat_cache=enc_feat_map, feat_idx=0) + spatial_sharding = _spatial_sharding(self.mesh, out_0.shape) if spatial_sharding is not None: out_0 = jax.lax.with_sharding_constraint(out_0, spatial_sharding) @@ -1280,11 +1290,7 @@ def _decode( x = self.post_quant_conv(z) dec_feat_map = feat_cache._feat_map - spatial_sharding = ( - NamedSharding(self.mesh, P("redundant", None, None, "vae_spatial", None)) - if self.mesh is not None and "vae_spatial" in self.mesh.shape - else None - ) + spatial_sharding = _spatial_sharding(self.mesh, x.shape) # First chunk (i=0) with jax.named_scope("AutoencoderKLWan_decode_chunk_0"):