Improvement in WAN VAE Performance

src/maxdiffusion/generate_wan.py

@@ -286,6 +286,8 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
           # would silently hit stale binaries.
           "flash_block_sizes": str(config.flash_block_sizes),
           "mesh_shape": str(pipeline.mesh.shape),
+          "vae_spatial": str(config.vae_spatial),
+          "vae_decode_chunk": str(config.vae_decode_chunk),
           "weights_dtype": str(config.weights_dtype),
           "activations_dtype": str(config.activations_dtype),
           "scan_layers": str(config.scan_layers),

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -23,6 +23,7 @@
 from jax import tree_util
 from flax import nnx
 from ...configuration_utils import ConfigMixin
+from ... import max_logging
 from ..modeling_flax_utils import FlaxModelMixin, get_activation
 from ... import common_types
 from ..vae_flax import (
@@ -67,6 +68,29 @@ def __eq__(self, other):
 tree_util.register_pytree_node(RepSentinel, lambda x: ((), None), lambda _, __: RepSentinel())
 
 
+def _with_sharding_constraint(x, sharding):
+  if sharding is not None and hasattr(x, "shape"):
+    if hasattr(sharding, "mesh") and hasattr(sharding, "spec") and sharding.mesh is not None:
+      mesh = sharding.mesh
+      spec = sharding.spec
+      # Guard against rank mismatch when mapping over heterogenous PyTree caches
+      if len(spec) != x.ndim:
+        return x
+      for axis_idx, axis_names in enumerate(spec):
+        if axis_names is not None:
+          if not isinstance(axis_names, tuple):
+            axis_names = (axis_names,)
+          mesh_axis_size = 1
+          for axis_name in axis_names:
+            if axis_name in mesh.shape:
+              mesh_axis_size *= mesh.shape[axis_name]
+          if x.shape[axis_idx] % mesh_axis_size != 0:
+            max_logging.log(f"Warning: Sharding mismatch. Shape {x.shape} at axis {axis_idx} does not divide evenly by mesh_axis_size {mesh_axis_size}. Skipping constraint.")
+            return x
+    return jax.lax.with_sharding_constraint(x, sharding)
+  return x
+
+
 class WanCausalConv3d(nnx.Module):
 
   def __init__(
@@ -99,9 +123,9 @@ def __init__(
     self.mesh = mesh
 
     # Weight sharding (Kernel is sharded along output channels)
-    num_fsdp_devices = mesh.shape["vae_spatial"]
+    num_fsdp_devices = mesh.shape["vae_spatial"] if mesh is not None and "vae_spatial" in mesh.shape else 1
     kernel_sharding = (None, None, None, None, None)
-    if out_channels % num_fsdp_devices == 0:
+    if num_fsdp_devices > 1 and out_channels % num_fsdp_devices == 0:
       kernel_sharding = (None, None, None, None, "vae_spatial")
 
     self.conv = nnx.Conv(
@@ -119,8 +143,9 @@ def __init__(
     )
 
   def __call__(self, x: jax.Array, cache_x: Optional[jax.Array] = None, idx=-1) -> jax.Array:
-    spatial_sharding = NamedSharding(self.mesh, P("redundant", None, None, "vae_spatial", None))
-    x = jax.lax.with_sharding_constraint(x, spatial_sharding)
+    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))
+      x = _with_sharding_constraint(x, spatial_sharding)
 
     current_padding = list(self._causal_padding)
     padding_needed = self._depth_padding_before
@@ -198,8 +223,16 @@ def __call__(self, x: jax.Array) -> jax.Array:
     n, h, w, c = in_shape
     target_h = int(h * self.scale_factor[0])
     target_w = int(w * self.scale_factor[1])
-    out = jax.image.resize(x.astype(jnp.float32), (n, target_h, target_w, c), method=self.method)
-    return out.astype(input_dtype)
+    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:
+      if self.method == "nearest":
+        max_logging.log(f"Warning: WanUpsample2D nearest method requested but scale_factor {self.scale_factor} is not integer. Falling back to jax.image.resize.")
+      out = jax.image.resize(x.astype(jnp.float32), (n, target_h, target_w, c), method=self.method)
+      out = out.astype(input_dtype)
+    return out
 
 
 class Identity(nnx.Module):
@@ -225,14 +258,16 @@ def __init__(
       weights_dtype: jnp.dtype = jnp.float32,
       precision: jax.lax.Precision = None,
   ):
+    rank = len(kernel_size) if isinstance(kernel_size, (tuple, list)) else 2
+    kernel_sharding = (None,) * (rank + 2)
     self.conv = nnx.Conv(
         dim,
         dim,
         kernel_size=kernel_size,
         strides=stride,
         use_bias=True,
         rngs=rngs,
-        kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), (None, None, None, None)),
+        kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), kernel_sharding),
         dtype=dtype,
         param_dtype=weights_dtype,
         precision=precision,
@@ -1131,7 +1166,6 @@ def __init__(
     )
     self.mesh = mesh
 
-  @nnx.jit
   def _encode(self, x: jax.Array, feat_cache: AutoencoderKLWanCache):
     feat_cache.init_cache()
     if x.shape[-1] != 3:
@@ -1151,7 +1185,11 @@ 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))
+    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
@@ -1162,7 +1200,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)
-      out_0 = jax.lax.with_sharding_constraint(out_0, spatial_sharding)
+      out_0 = _with_sharding_constraint(out_0, spatial_sharding)
 
     if iter_ <= 1:
       return finalize(out_0, enc_feat_map)
@@ -1172,11 +1210,11 @@ def finalize(out, enc_feat_map):
     with jax.named_scope("AutoencoderKLWan_encode_chunk_1"):
       chunk_1 = x[:, 1 : (1 + CHUNK_SIZE), ...]
       out_1, enc_feat_map, _ = self.encoder(chunk_1, feat_cache=enc_feat_map, feat_idx=0)
-      out_1 = jax.lax.with_sharding_constraint(out_1, spatial_sharding)
+      out_1 = _with_sharding_constraint(out_1, spatial_sharding)
 
     if iter_ <= 2:
       out = jnp.concatenate([out_0, out_1], axis=1)
-      out = jax.lax.with_sharding_constraint(out, spatial_sharding)
+      out = _with_sharding_constraint(out, spatial_sharding)
       return finalize(out, enc_feat_map)
 
     # Prepare the remaining chunks to be scanned over
@@ -1209,10 +1247,11 @@ def finalize(out, enc_feat_map):
     def scan_fn(carry, chunk):
       current_feat_map = carry
       local_encoder = nnx.merge(graphdef, state)
+      chunk = _with_sharding_constraint(chunk, spatial_sharding)
       out_chunk, next_feat_map, _ = local_encoder(chunk, feat_cache=current_feat_map, feat_idx=0)
-      out_chunk = jax.lax.with_sharding_constraint(out_chunk, spatial_sharding)
+      out_chunk = _with_sharding_constraint(out_chunk, spatial_sharding)
       next_feat_map = jax.tree_util.tree_map(
-          lambda x: jax.lax.with_sharding_constraint(x, spatial_sharding) if isinstance(x, jax.Array) else x, next_feat_map
+          lambda x: _with_sharding_constraint(x, spatial_sharding) if hasattr(x, "shape") else x, next_feat_map
       )
       return next_feat_map, out_chunk
 
@@ -1225,7 +1264,7 @@ def scan_fn(carry, chunk):
     out_rest = out_rest[:, : T_rest // self.temporal_downsample_factor, ...]
 
     out = jnp.concatenate([out_0, out_1, out_rest], axis=1)
-    out = jax.lax.with_sharding_constraint(out, spatial_sharding)
+    out = _with_sharding_constraint(out, spatial_sharding)
     return finalize(out, enc_feat_map)
 
   @jax.named_scope("AutoencoderKLWan_encode")
@@ -1239,7 +1278,6 @@ def encode(
       return (posterior,)
     return FlaxAutoencoderKLOutput(latent_dist=posterior)
 
-  @nnx.jit
   def _decode(
       self, z: jax.Array, feat_cache: AutoencoderKLWanCache, return_dict: bool = True
   ) -> Union[FlaxDecoderOutput, jax.Array]:
@@ -1249,20 +1287,24 @@ 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))
+    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
+    )
 
     # First chunk (i=0)
     with jax.named_scope("AutoencoderKLWan_decode_chunk_0"):
-      chunk_in_0 = jax.lax.with_sharding_constraint(x[:, 0:1, ...], spatial_sharding)
+      chunk_in_0 = _with_sharding_constraint(x[:, 0:1, ...], spatial_sharding)
       out_0, dec_feat_map, _ = self.decoder(chunk_in_0, feat_cache=dec_feat_map, feat_idx=0)
-      out_0 = jax.lax.with_sharding_constraint(out_0, spatial_sharding)
+      out_0 = _with_sharding_constraint(out_0, spatial_sharding)
 
     if iter_ > 1:
       # Run chunk 1 outside scan to properly form the cache shape
       with jax.named_scope("AutoencoderKLWan_decode_chunk_1"):
-        chunk_in_1 = jax.lax.with_sharding_constraint(x[:, 1:2, ...], spatial_sharding)
+        chunk_in_1 = _with_sharding_constraint(x[:, 1:2, ...], spatial_sharding)
         out_chunk_1, dec_feat_map, _ = self.decoder(chunk_in_1, feat_cache=dec_feat_map, feat_idx=0)
-        out_chunk_1 = jax.lax.with_sharding_constraint(out_chunk_1, spatial_sharding)
+        out_chunk_1 = _with_sharding_constraint(out_chunk_1, spatial_sharding)
 
       out_1 = out_chunk_1
       out_list = [out_0, out_1]
@@ -1297,11 +1339,11 @@ def _decode(
         def scan_fn(carry, chunk_in):
           current_feat_map = carry
           local_decoder = nnx.merge(graphdef, state)
-          chunk_in = jax.lax.with_sharding_constraint(chunk_in, spatial_sharding)
+          chunk_in = _with_sharding_constraint(chunk_in, spatial_sharding)
           out_chunk, next_feat_map, _ = local_decoder(chunk_in, feat_cache=current_feat_map, feat_idx=0)
-          out_chunk = jax.lax.with_sharding_constraint(out_chunk, spatial_sharding)
+          out_chunk = _with_sharding_constraint(out_chunk, spatial_sharding)
           next_feat_map = jax.tree_util.tree_map(
-              lambda x: jax.lax.with_sharding_constraint(x, spatial_sharding) if isinstance(x, jax.Array) else x,
+              lambda x: _with_sharding_constraint(x, spatial_sharding) if hasattr(x, "shape") else x,
               next_feat_map,
           )
           return next_feat_map, out_chunk
@@ -1314,7 +1356,7 @@ def scan_fn(carry, chunk_in):
         out_list.append(out_rest)
 
       out = jnp.concatenate(out_list, axis=1)
-      out = jax.lax.with_sharding_constraint(out, spatial_sharding)
+      out = _with_sharding_constraint(out, spatial_sharding)
     else:
       out = out_0
 

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -474,6 +474,38 @@ def __init__(
     # repeated serving requests) skip the ~10s/call CPU text encoder.
     self._prompt_embeds_cache = {}
 
+  def check_inputs(
+      self,
+      prompt: Union[str, List[str]] = None,
+      negative_prompt: Optional[Union[str, List[str]]] = None,
+      height: int = 480,
+      width: int = 832,
+      prompt_embeds: Optional[jax.Array] = None,
+      negative_prompt_embeds: Optional[jax.Array] = None,
+      **kwargs,
+  ):
+    """Validate user-facing pipeline inputs and shape contracts."""
+    if prompt is not None and prompt_embeds is not None:
+      raise ValueError(
+          f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+          " only forward one of the two."
+      )
+    elif negative_prompt is not None and negative_prompt_embeds is not None:
+      raise ValueError(
+          f"Cannot forward both `negative_prompt`: {negative_prompt} and"
+          f" `negative_prompt_embeds`: {negative_prompt_embeds}. Please make sure to"
+          " only forward one of the two."
+      )
+
+    mesh = getattr(self, "vae_mesh", getattr(self, "mesh", None))
+    if mesh is not None and hasattr(mesh, "shape"):
+      vae_spatial = mesh.shape.get("vae_spatial", 1)
+      if vae_spatial > 1 and (width // 8) % vae_spatial != 0:
+        max_logging.log(
+            f"Warning: Latent width is not divisible by vae_spatial mesh axis ({vae_spatial})."
+            " VAE spatial sharding will be partially bypassed."
+        )
+
   @classmethod
   def load_text_encoder(cls, config: HyperParameters):
     text_encoder_dtype = getattr(config, "text_encoder_dtype", "float32")
@@ -907,8 +939,10 @@ def _decode_latents_to_video(self, latents: jax.Array, trace: Optional[dict] = N
     if trace is not None:
       trace["vae_decode_tpu"] = time.perf_counter() - t_vae_tpu_start
 
-    video = jax.experimental.multihost_utils.process_allgather(video, tiled=True)
-    video = np.array(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)
     return video
 
   @classmethod
@@ -1231,6 +1265,14 @@ def _prepare_model_inputs(
       prompt_embeds: jax.Array = None,
       negative_prompt_embeds: jax.Array = None,
   ):
+    self.check_inputs(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        height=height,
+        width=width,
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=negative_prompt_embeds,
+    )
     if max_sequence_length is None:
       max_sequence_length = getattr(self.config, "max_sequence_length", 512)
 
@@ -1315,7 +1357,6 @@ def __call__(self, **kwargs):
     aot_cache.cached_jit,
     static_argnames=(
         "do_classifier_free_guidance",
-        "guidance_scale",
         "return_residual",
         "skip_blocks",
     ),
@@ -1337,6 +1378,8 @@ def transformer_forward_pass(
     rotary_emb=None,
     encoder_attention_mask=None,
 ):
+  if do_classifier_free_guidance and latents.shape[0] != prompt_embeds.shape[0]:
+    latents = jnp.concatenate([latents, latents], axis=0)
   wan_transformer = nnx.merge(graphdef, sharded_state, rest_of_state)
   outputs = wan_transformer(
       hidden_states=latents,
@@ -1362,11 +1405,9 @@ def transformer_forward_pass(
     noise_uncond = noise_pred[bsz:]  # Second half = unconditional
     noise_pred = noise_uncond + guidance_scale * (noise_cond - noise_uncond)
 
-    latents = latents[:bsz]
-
   if return_residual:
-    return noise_pred, latents, residual_x
-  return noise_pred, latents
+    return noise_pred, residual_x
+  return noise_pred
 
 
 @aot_cache.cached_jit
@@ -1389,10 +1430,14 @@ def vae_decode_pass(graphdef, state, rest_of_state, latents):
   video = wan_vae.decode(latents, AutoencoderKLWanCache(wan_vae), return_dict=False)[0]
   video = (video / 2.0) + 0.5
   video = jnp.clip(video, 0.0, 1.0)
-  return (video * 255.0).astype(jnp.uint8)
+  video = (video * 255.0).astype(jnp.uint8)
+  if wan_vae.mesh is not None:
+    replicated_sharding = NamedSharding(wan_vae.mesh, P())
+    video = jax.lax.with_sharding_constraint(video, replicated_sharding)
+  return video
 
 
-@partial(aot_cache.cached_jit, static_argnames=("guidance_scale",))
+@aot_cache.cached_jit
 def transformer_forward_pass_full_cfg(
     graphdef,
     sharded_state,
@@ -1433,7 +1478,7 @@ def transformer_forward_pass_full_cfg(
   return noise_pred_merged, noise_cond, noise_uncond
 
 
-@partial(aot_cache.cached_jit, static_argnames=("guidance_scale",))
+@aot_cache.cached_jit
 def transformer_forward_pass_cfg_cache(
     graphdef,
     sharded_state,