From a838e53bdfffb8480651828b14f282cc2f141724 Mon Sep 17 00:00:00 2001
From: Toshi Pahadia <toshipahadia@google.com>
Date: Thu, 20 Aug 2026 10:50:08 +0530
Subject: [PATCH] Optimize WAN Pipeline and VAE Inference

- 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
---
 src/maxdiffusion/generate_wan.py              |  2 +
 .../models/wan/autoencoder_kl_wan.py          | 92 ++++++++++++++-----
 .../pipelines/wan/wan_pipeline.py             | 65 +++++++++++--
 .../pipelines/wan/wan_pipeline_2_1.py         | 14 +--
 .../pipelines/wan/wan_pipeline_2_2.py         | 69 +++++++-------
 .../pipelines/wan/wan_pipeline_animate.py     |  9 ++
 .../pipelines/wan/wan_pipeline_i2v_2p1.py     |  7 +-
 .../pipelines/wan/wan_pipeline_i2v_2p2.py     | 62 +++++++------
 .../pipelines/wan/wan_vace_pipeline_2_1.py    | 39 +++++---
 src/maxdiffusion/tests/models_import_test.py  |  6 +-
 src/maxdiffusion/tests/profiler_test.py       |  8 +-
 src/maxdiffusion/tests/wan/wan_vae_test.py    |  6 +-
 12 files changed, 240 insertions(+), 139 deletions(-)

diff --git a/src/maxdiffusion/generate_wan.py b/src/maxdiffusion/generate_wan.py
index 01199a303..b80f46cd4 100644
--- a/src/maxdiffusion/generate_wan.py
+++ b/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),
diff --git a/src/maxdiffusion/models/wan/autoencoder_kl_wan.py b/src/maxdiffusion/models/wan/autoencoder_kl_wan.py
index 89f2647b5..06fe4bce4 100644
--- a/src/maxdiffusion/models/wan/autoencoder_kl_wan.py
+++ b/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,6 +258,8 @@ 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,
@@ -232,7 +267,7 @@ def __init__(
         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
 
diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline.py b/src/maxdiffusion/pipelines/wan/wan_pipeline.py
index dc3b3fbc4..80e150ce6 100644
--- a/src/maxdiffusion/pipelines/wan/wan_pipeline.py
+++ b/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,
diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py
index a963d3498..a6b97cc10 100644
--- a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py
+++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py
@@ -359,15 +359,15 @@ def scan_body(carry, t):
       current_latents, current_scheduler_state = carry
 
       if do_cfg:
-        latents_doubled = jnp.concatenate([current_latents] * 2)
         timestep = jnp.broadcast_to(t, bsz * 2)
-        noise_pred, _, _ = transformer_forward_pass_full_cfg(
+        noise_pred = transformer_forward_pass(
             graphdef,
             sharded_state,
             rest_of_state,
-            latents_doubled,
+            current_latents,
             timestep,
             prompt_embeds_combined,
+            do_classifier_free_guidance=True,
             guidance_scale=guidance_scale,
             kv_cache=kv_cache,
             rotary_emb=rotary_emb,
@@ -375,7 +375,7 @@ def scan_body(carry, t):
         )
       else:
         timestep = jnp.broadcast_to(t, bsz)
-        noise_pred, _ = transformer_forward_pass(
+        noise_pred = transformer_forward_pass(
             graphdef,
             sharded_state,
             rest_of_state,
@@ -422,11 +422,11 @@ def scan_body(carry, t):
           skip_warmup,
       )
 
-      noise_pred, latents, residual_x_cur = transformer_forward_pass(
+      noise_pred, residual_x_cur = transformer_forward_pass(
           graphdef,
           sharded_state,
           rest_of_state,
-          jnp.concatenate([latents] * 2) if do_cfg else latents,
+          latents,
           timestep,
           prompt_embeds_combined if do_cfg else prompt_cond_embeds,
           do_classifier_free_guidance=do_cfg,
@@ -489,7 +489,7 @@ def scan_body(carry, t):
 
       else:
         timestep = jnp.broadcast_to(t, bsz)
-        noise_pred, latents = transformer_forward_pass(
+        noise_pred = transformer_forward_pass(
             graphdef,
             sharded_state,
             rest_of_state,
diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py
index 5aac990a7..29546ff9a 100644
--- a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py
+++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py
@@ -435,13 +435,12 @@ def run_inference_2_2(
           use_magcache=(not force_compute),
       )
 
-      latents_doubled = jnp.concatenate([latents] * 2)
       timestep = jnp.broadcast_to(t, bsz * 2)
-      noise_pred, _, residual_x_cur = transformer_forward_pass(
+      noise_pred, residual_x_cur = transformer_forward_pass(
           graphdef,
           state,
           rest,
-          latents_doubled,
+          latents,
           timestep,
           prompt_embeds_combined,
           do_classifier_free_guidance=True,
@@ -566,22 +565,23 @@ def run_inference_2_2(
         continue
 
       # Accumulate deltas since last full compute
-      dx_norm = float(jnp.sqrt(jnp.mean((latents - ref_latent) ** 2)))
-      dt = abs(t_float - ref_timestep)
-      accum_dx += dx_norm
-      accum_dt += dt
+      dx_norm = jnp.sqrt(jnp.mean((latents - ref_latent) ** 2))
+      dt = jnp.abs(t_float - ref_timestep)
+      accum_dx = accum_dx + dx_norm
+      accum_dt = accum_dt + dt
 
       # Sensitivity score (Eq. 9)
       score = alpha_x * accum_dx + alpha_t * accum_dt
+      should_reuse = (score <= sen_epsilon) & (reuse_count < max_reuse)
 
-      if score <= sen_epsilon and reuse_count < max_reuse:
-        noise_pred = ref_noise_pred
-        reuse_count += 1
-        cache_count += 1
-      else:
-        latents_doubled = jnp.concatenate([latents] * 2)
-        timestep = jnp.broadcast_to(t, bsz * 2)
-        noise_pred, _, _ = transformer_forward_pass_full_cfg(
+      latents_doubled = jnp.concatenate([latents] * 2)
+      timestep = jnp.broadcast_to(t, bsz * 2)
+
+      def reuse_fn():
+        return ref_noise_pred
+
+      def compute_fn():
+        out, _, _ = transformer_forward_pass_full_cfg(
             graphdef,
             state,
             rest,
@@ -593,12 +593,16 @@ def run_inference_2_2(
             rotary_emb=rotary_emb,
             encoder_attention_mask=encoder_attention_mask,
         )
-        ref_noise_pred = noise_pred
-        ref_latent = latents
-        ref_timestep = t_float
-        accum_dx = 0.0
-        accum_dt = 0.0
-        reuse_count = 0
+        return out
+
+      noise_pred = jax.lax.cond(should_reuse, reuse_fn, compute_fn)
+      ref_noise_pred = jnp.where(should_reuse, ref_noise_pred, noise_pred)
+      ref_latent = jnp.where(should_reuse, ref_latent, latents)
+      ref_timestep = jnp.where(should_reuse, ref_timestep, t_float)
+      accum_dx = jnp.where(should_reuse, accum_dx, 0.0)
+      accum_dt = jnp.where(should_reuse, accum_dt, 0.0)
+      reuse_count = jnp.where(should_reuse, reuse_count + 1, 0)
+      cache_count = jnp.where(should_reuse, cache_count + 1, cache_count)
 
       latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
 
@@ -813,22 +817,17 @@ def low_noise_branch(operands):
     def scan_body(carry, t):
       current_latents, current_scheduler_state = carry
 
-      if do_classifier_free_guidance:
-        model_latents = jnp.concatenate([current_latents] * 2)
-      else:
-        model_latents = current_latents
-
-      timestep = jnp.broadcast_to(t, model_latents.shape[0])
+      timestep = jnp.broadcast_to(t, (bsz * 2 if do_classifier_free_guidance else bsz,))
       use_high_noise = jnp.greater_equal(t, boundary)
 
-      noise_pred, latents_out = jax.lax.cond(
+      noise_pred = jax.lax.cond(
           use_high_noise,
           high_noise_branch,
           low_noise_branch,
           (
-              model_latents,
+              current_latents,
               timestep,
-              prompt_embeds_combined,
+              prompt_embeds_combined if do_classifier_free_guidance else prompt_embeds,
               kv_cache_high,
               kv_cache_low,
               rotary_emb,
@@ -838,7 +837,7 @@ def scan_body(carry, t):
       )
 
       new_latents, new_scheduler_state = scheduler.step(
-          current_scheduler_state, noise_pred, t, latents_out, return_dict=False
+          current_scheduler_state, noise_pred, t, current_latents, return_dict=False
       )
 
       return (new_latents, new_scheduler_state), None
@@ -917,15 +916,15 @@ def scan_body(carry, t):
       encoder_attention_mask = encoder_attention_mask_low
 
     if do_classifier_free_guidance:
-      latents_doubled = jnp.concatenate([latents] * 2)
       timestep = jnp.broadcast_to(t, bsz * 2)
-      noise_pred, _, _ = transformer_forward_pass_full_cfg(
+      noise_pred = transformer_forward_pass(
           graphdef,
           state,
           rest,
-          latents_doubled,
+          latents,
           timestep,
           prompt_embeds_combined,
+          do_classifier_free_guidance=True,
           guidance_scale=guidance_scale,
           kv_cache=kv_cache,
           rotary_emb=rotary_emb,
@@ -933,7 +932,7 @@ def scan_body(carry, t):
       )
     else:
       timestep = jnp.broadcast_to(t, bsz)
-      noise_pred, latents = transformer_forward_pass(
+      noise_pred = transformer_forward_pass(
           graphdef,
           state,
           rest,
diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_animate.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_animate.py
index f01d0ea4d..6a926f9ea 100644
--- a/src/maxdiffusion/pipelines/wan/wan_pipeline_animate.py
+++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_animate.py
@@ -412,6 +412,15 @@ def check_inputs(
           f" {type(prev_segment_conditioning_frames)} and value is {prev_segment_conditioning_frames}"
       )
 
+    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."
+        )
+
   @staticmethod
   def pad_video_frames(frames: list, num_target_frames: int) -> list:
     """Pad *frames* to *num_target_frames* using a reflect-like strategy.
diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py
index 5fc35d8d7..a3c1ff514 100644
--- a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py
+++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py
@@ -389,7 +389,7 @@ def scan_body(carry, t):
       latent_model_input = jnp.concatenate([latents_input, condition_combined], axis=1)
       timestep = jnp.broadcast_to(t, latents_input.shape[0])
 
-      outputs = transformer_forward_pass(
+      noise_pred = transformer_forward_pass(
           graphdef,
           sharded_state,
           rest_of_state,
@@ -406,7 +406,6 @@ def scan_body(carry, t):
           rotary_emb=rotary_emb,
           encoder_attention_mask=encoder_attention_mask,
       )
-      noise_pred, _ = outputs
 
       noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1))
       new_latents, new_scheduler_state = scheduler.step(
@@ -467,11 +466,11 @@ def scan_body(carry, t):
         encoder_attention_mask=encoder_attention_mask,
     )
     if use_magcache and do_cfg:
-      noise_pred, _, residual_x_cur = outputs
+      noise_pred, residual_x_cur = outputs
       if not skip_blocks:
         cached_residual = residual_x_cur
     else:
-      noise_pred, _ = outputs
+      noise_pred = outputs
 
     noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1))
     latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents, return_dict=False)
diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py
index 3436fe8a7..7e7997030 100644
--- a/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py
+++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py
@@ -547,7 +547,7 @@ def run_inference_2_2_i2v(
       latent_model_input = jnp.concatenate([latents_doubled, condition_doubled], axis=-1)
       latent_model_input = jnp.transpose(latent_model_input, (0, 4, 1, 2, 3))
       timestep = jnp.broadcast_to(t, bsz * 2)
-      noise_pred, _, residual_x_cur = transformer_forward_pass(
+      noise_pred, residual_x_cur = transformer_forward_pass(
           graphdef,
           state,
           rest,
@@ -646,13 +646,14 @@ def run_inference_2_2_i2v(
         latents_doubled = jnp.transpose(latents_doubled, (0, 4, 1, 2, 3))
         latent_model_input = jnp.concatenate([latents_doubled, condition_doubled], axis=1)
         timestep = jnp.broadcast_to(t, bsz * 2)
-        noise_pred, _, _ = transformer_forward_pass_full_cfg(
+        noise_pred = transformer_forward_pass(
             graphdef,
             state,
             rest,
             latent_model_input,
             timestep,
             prompt_embeds_combined,
+            do_classifier_free_guidance=True,
             guidance_scale=guidance_scale,
             encoder_hidden_states_image=image_embeds_combined,
             kv_cache=kv_cache,
@@ -669,42 +670,47 @@ def run_inference_2_2_i2v(
         latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
         continue
 
-      dx_norm = float(jnp.sqrt(jnp.mean((latents - ref_latent) ** 2)))
-      dt = abs(t_float - ref_timestep)
-      accum_dx += dx_norm
-      accum_dt += dt
+      dx_norm = jnp.sqrt(jnp.mean((latents - ref_latent) ** 2))
+      dt = jnp.abs(t_float - ref_timestep)
+      accum_dx = accum_dx + dx_norm
+      accum_dt = accum_dt + dt
 
       score = alpha_x * accum_dx + alpha_t * accum_dt
+      should_reuse = (score <= sen_epsilon) & (reuse_count < max_reuse)
 
-      if score <= sen_epsilon and reuse_count < max_reuse:
-        noise_pred = ref_noise_pred
-        reuse_count += 1
-        cache_count += 1
-      else:
-        latents_doubled = jnp.concatenate([latents, latents], axis=0)
-        latents_doubled = jnp.transpose(latents_doubled, (0, 4, 1, 2, 3))
-        latent_model_input = jnp.concatenate([latents_doubled, condition_doubled], axis=1)
-        timestep = jnp.broadcast_to(t, bsz * 2)
-        noise_pred, _, _ = transformer_forward_pass_full_cfg(
+      latents_doubled = jnp.concatenate([latents, latents], axis=0)
+      latents_doubled = jnp.transpose(latents_doubled, (0, 4, 1, 2, 3))
+      latent_model_input = jnp.concatenate([latents_doubled, condition_doubled], axis=1)
+      timestep = jnp.broadcast_to(t, bsz * 2)
+
+      def reuse_fn():
+        return ref_noise_pred
+
+      def compute_fn():
+        out = transformer_forward_pass(
             graphdef,
             state,
             rest,
             latent_model_input,
             timestep,
             prompt_embeds_combined,
+            do_classifier_free_guidance=True,
             guidance_scale=guidance_scale,
             encoder_hidden_states_image=image_embeds_combined,
             kv_cache=kv_cache,
             rotary_emb=rotary_emb,
             encoder_attention_mask=encoder_attention_mask,
         )
-        noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1))
-        ref_noise_pred = noise_pred
-        ref_latent = latents
-        ref_timestep = t_float
-        accum_dx = 0.0
-        accum_dt = 0.0
-        reuse_count = 0
+        return jnp.transpose(out, (0, 2, 3, 4, 1))
+
+      noise_pred = jax.lax.cond(should_reuse, reuse_fn, compute_fn)
+      ref_noise_pred = jnp.where(should_reuse, ref_noise_pred, noise_pred)
+      ref_latent = jnp.where(should_reuse, ref_latent, latents)
+      ref_timestep = jnp.where(should_reuse, ref_timestep, t_float)
+      accum_dx = jnp.where(should_reuse, accum_dx, 0.0)
+      accum_dt = jnp.where(should_reuse, accum_dt, 0.0)
+      reuse_count = jnp.where(should_reuse, reuse_count + 1, 0)
+      cache_count = jnp.where(should_reuse, cache_count + 1, cache_count)
 
       latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
 
@@ -867,7 +873,7 @@ def high_noise_branch(operands):
         mask_high,
         _,
     ) = operands
-    noise_pred, latents_out = transformer_forward_pass(
+    return transformer_forward_pass(
         high_noise_graphdef,
         high_noise_state,
         high_noise_rest,
@@ -881,7 +887,6 @@ def high_noise_branch(operands):
         rotary_emb=r_emb,
         encoder_attention_mask=mask_high,
     )
-    return noise_pred, latents_out
 
   def low_noise_branch(operands):
     (
@@ -895,7 +900,7 @@ def low_noise_branch(operands):
         _,
         mask_low,
     ) = operands
-    noise_pred, latents_out = transformer_forward_pass(
+    return transformer_forward_pass(
         low_noise_graphdef,
         low_noise_state,
         low_noise_rest,
@@ -909,7 +914,6 @@ def low_noise_branch(operands):
         rotary_emb=r_emb,
         encoder_attention_mask=mask_low,
     )
-    return noise_pred, latents_out
 
   if do_classifier_free_guidance:
     condition = jnp.concatenate([condition] * 2)
@@ -939,7 +943,7 @@ def scan_body(carry, t):
       timestep = jnp.broadcast_to(t, latents_input.shape[0])
 
       use_high_noise = jnp.greater_equal(t, boundary)
-      noise_pred, _ = jax.lax.cond(
+      noise_pred = jax.lax.cond(
           use_high_noise,
           high_noise_branch,
           low_noise_branch,
@@ -987,7 +991,7 @@ def scan_body(carry, t):
     # tracing both 14B branches per step and keeps the AOT cache usable.
     use_high_noise = bool(np.asarray(scheduler_state.timesteps)[step] >= np.asarray(boundary))
     branch = high_noise_branch if use_high_noise else low_noise_branch
-    noise_pred, _ = branch((
+    noise_pred = branch((
         latent_model_input,
         timestep,
         prompt_embeds_combined,
diff --git a/src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py b/src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py
index e2c6a91d0..8090f7590 100644
--- a/src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py
+++ b/src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py
@@ -30,7 +30,7 @@
 from ...models.wan.transformers.transformer_wan_vace import WanVACEModel
 from ...schedulers.scheduling_unipc_multistep_flax import FlaxUniPCMultistepScheduler
 from ...models.modeling_flax_pytorch_utils import torch2jax
-from .wan_pipeline import _final_param_dtype, cast_with_exclusion, converted_weights_cache_dir, put_params_into_state
+from .wan_pipeline import _final_param_dtype, cast_with_exclusion, converted_weights_cache_dir, put_params_into_state, vae_encode_pass
 from .wan_pipeline_2_1 import WanPipeline2_1
 import torch
 import PIL
@@ -459,6 +459,15 @@ def check_inputs(
     elif mask is not None:
       raise ValueError("`mask` can only be passed if `video` is passed as well.")
 
+    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."
+        )
+
   def __call__(
       self,
       video: Optional[List[PipelineImageInput]] = None,
@@ -681,8 +690,9 @@ def prepare_video_latents(
       inactive = video * (1 - mask)
       reactive = video * mask
       with self.vae_mesh, nn_partitioning.axis_rules(self.vae_logical_axis_rules):
-        inactive = retrieve_latents(self.vae.encode(inactive, self.vae_cache), rngs=rngs, sample_mode="argmax")
-        reactive = retrieve_latents(self.vae.encode(reactive, self.vae_cache), rngs=rngs, sample_mode="argmax")
+        graphdef, state, rest_of_state = nnx.split(self.vae, nnx.Param, ...)
+        inactive = vae_encode_pass(graphdef, state, rest_of_state, inactive)
+        reactive = vae_encode_pass(graphdef, state, rest_of_state, reactive)
       inactive = ((inactive.astype(jnp.float32) - latents_mean) * latents_std).astype(vae_dtype)
       reactive = ((reactive.astype(jnp.float32) - latents_mean) * latents_std).astype(vae_dtype)
 
@@ -697,9 +707,8 @@ def prepare_video_latents(
         reference_image = reference_image[None, None, :, :, :]  # [1, 1, H, W, C]
 
         with self.vae_mesh, nn_partitioning.axis_rules(self.vae_logical_axis_rules):
-          reference_latent = retrieve_latents(
-              self.vae.encode(reference_image, feat_cache=self.vae_cache), rngs=None, sample_mode="argmax"
-          )
+          graphdef, state, rest_of_state = nnx.split(self.vae, nnx.Param, ...)
+          reference_latent = vae_encode_pass(graphdef, state, rest_of_state, reference_image)
 
         reference_latent = ((reference_latent.astype(jnp.float32) - latents_mean) * latents_std).astype(vae_dtype)
 
@@ -712,7 +721,7 @@ def prepare_video_latents(
     return jnp.stack(latent_list)
 
 
-@partial(aot_cache.cached_jit, static_argnames=("do_classifier_free_guidance", "guidance_scale"))
+@partial(aot_cache.cached_jit, static_argnames=("do_classifier_free_guidance",))
 def transformer_forward_pass(
     graphdef: nnx.graph.GraphDef,
     sharded_state: nnx.State,
@@ -728,6 +737,8 @@ def transformer_forward_pass(
     encoder_attention_mask=None,
 ):
   """Performs a forward pass on the transformer."""
+  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)
   noise_pred = wan_transformer(
       hidden_states=latents,
@@ -743,9 +754,8 @@ def transformer_forward_pass(
     noise_uncond = noise_pred[bsz:]
     noise_pred = noise_pred[:bsz]
     noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond)
-    latents = latents[:bsz]
 
-  return noise_pred, latents
+  return noise_pred
 
 
 def run_inference(
@@ -774,13 +784,13 @@ def run_inference(
   if use_kv_cache:
     kv_cache, encoder_attention_mask = transformer_obj.compute_kv_cache(prompt_embeds)
 
+  bsz = latents.shape[0] * 2 if do_classifier_free_guidance else latents.shape[0]
+  timesteps = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)
   for step in range(num_inference_steps):
-    t = jnp.array(scheduler_state.timesteps, dtype=jnp.int32)[step]
-    if do_classifier_free_guidance:
-      latents = jnp.concatenate([latents] * 2)
-    timestep = jnp.broadcast_to(t, latents.shape[0])
+    t = timesteps[step]
+    timestep = jnp.broadcast_to(t, (bsz,))
 
-    noise_pred, latents = transformer_forward_pass(
+    noise_pred = transformer_forward_pass(
         graphdef,
         sharded_state,
         rest_of_state,
@@ -796,4 +806,5 @@ def run_inference(
     )
 
     latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
+
   return latents
diff --git a/src/maxdiffusion/tests/models_import_test.py b/src/maxdiffusion/tests/models_import_test.py
index 3bd898cde..664345e63 100644
--- a/src/maxdiffusion/tests/models_import_test.py
+++ b/src/maxdiffusion/tests/models_import_test.py
@@ -50,8 +50,6 @@ def test_import_all_root_utilities(self):
         configuration_utils,
         max_logging,
         max_utils,
-        maxdiffusion_google,
-        maxdiffusion_google_hub,
         maxdiffusion_utils,
         multihost_dataloading,
         pyconfig,
@@ -64,8 +62,6 @@ def test_import_all_root_utilities(self):
     self.assertIsNotNone(configuration_utils)
     self.assertIsNotNone(max_logging)
     self.assertIsNotNone(max_utils)
-    self.assertIsNotNone(maxdiffusion_google)
-    self.assertIsNotNone(maxdiffusion_google_hub)
     self.assertIsNotNone(maxdiffusion_utils)
     self.assertIsNotNone(multihost_dataloading)
     self.assertIsNotNone(pyconfig)
@@ -74,7 +70,7 @@ def test_import_all_root_utilities(self):
 
   def test_mldiagnostics_import_and_usage(self):
     try:
-      from google_cloud_mldiagnostics import machinelearning_run, xprof
+      from google_cloud_mldiagnostics import machinelearning_run
       self.assertIsNotNone(machinelearning_run)
     except ImportError:
       pass
diff --git a/src/maxdiffusion/tests/profiler_test.py b/src/maxdiffusion/tests/profiler_test.py
index 763be4ca4..653339471 100644
--- a/src/maxdiffusion/tests/profiler_test.py
+++ b/src/maxdiffusion/tests/profiler_test.py
@@ -66,9 +66,7 @@ def test_ml_diagnostics_profiler(self, mock_process_index, mock_xprof, mock_ml_r
   @patch("maxdiffusion.max_utils.machinelearning_run")
   @patch("maxdiffusion.max_utils.xprof")
   @patch("jax.process_index", return_value=1)
-  def test_ml_diagnostics_profiler_non_master_host(
-      self, mock_process_index, mock_xprof, mock_ml_run
-  ):
+  def test_ml_diagnostics_profiler_non_master_host(self, mock_process_index, mock_xprof, mock_ml_run):
     """Tests that ML Diagnostics profiler is also enabled on non-master hosts (process_index != 0)."""
     config = MockConfig(
         enable_ml_diagnostics=True,
@@ -90,9 +88,7 @@ def test_ml_diagnostics_profiler_non_master_host(
   @patch("maxdiffusion.max_utils.machinelearning_run")
   @patch("maxdiffusion.max_utils.xprof")
   @patch("jax.process_index", return_value=0)
-  def test_both_profilers_enabled_prioritizes_mld(
-      self, mock_process_index, mock_xprof, mock_ml_run, mock_start_trace
-  ):
+  def test_both_profilers_enabled_prioritizes_mld(self, mock_process_index, mock_xprof, mock_ml_run, mock_start_trace):
     """Tests that when both ML Diagnostics and JAX profiler are enabled, ML Diagnostics is prioritized and JAX profiler is skipped."""
     config = MockConfig(
         enable_ml_diagnostics=True,
diff --git a/src/maxdiffusion/tests/wan/wan_vae_test.py b/src/maxdiffusion/tests/wan/wan_vae_test.py
index 04b8010bc..a1d52bb2c 100644
--- a/src/maxdiffusion/tests/wan/wan_vae_test.py
+++ b/src/maxdiffusion/tests/wan/wan_vae_test.py
@@ -226,11 +226,9 @@ def test_zero_padded_conv(self):
 
     with self.mesh, nn_partitioning.axis_rules(self.config.vae_logical_axis_rules):
       model = ZeroPaddedConv2D(dim=dim, rngs=rngs, kernel_size=(1, 3, 3), stride=(1, 2, 2))
-      dummy_input = jnp.ones(input_shape)
-      dummy_input = jnp.transpose(dummy_input, (0, 2, 3, 1))
+      dummy_input = jnp.ones((1, 1, 480, 720, 96))
       output = model(dummy_input)
-      output = jnp.transpose(output, (0, 3, 1, 2))
-      assert output.shape == (1, 96, 240, 360)
+      assert output.shape == (1, 1, 240, 360, 96)
 
   def test_wan_upsample(self):
     batch_size = 1