Fix VAE timing regression for large batch sizes

**Root Cause:**
In commit 7b28885, an optimization was added to prevent OOM errors for large batch sizes (`batch_size > 2`) by batch-sharding the latents and disabling sequential slicing. However, this logic used an `elif replicate_vae:` block, which caused the explicit replication of VAE weights to be entirely skipped for large batch sizes.

Without explicit weight replication, the XLA SPMD partitioner attempts to match the sharding of the input `latents` (which are batch-sharded) with the VAE decode computation. Because `vae.decode` involves fully-replicated noise injection and massive 3D convolutions, XLA heuristically decides to insert enormous amounts of cross-device communication (AllGather/AllReduce) to shard the weights or activations, ballooning the execution time from ~2.8s to ~68.5s for non-upsampled latents.

(Note: For upsampled latents, the memory layout generated by the JIT-compiled upsampler bypasses this XLA heuristic trap, allowing it to execute quickly in ~1.5s, which masked the issue).

**Fix:**
This PR decouples the batch-sharding of `latents` from the replication of VAE weights. It explicitly applies a full replication constraint `NamedSharding(mesh, P())` to the VAE weights in all cases where `replicate_vae` is True, even if `latents` are batch-sharded. This forces XLA into the optimal data-parallel compilation path, restoring the fast ~1.5s - ~2.8s execution time for all scenarios without risking the concatenation-related OOMs.

Author: mbohlool

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1818,10 +1818,9 @@ def convert_to_vel(lat, x0, sig):
     enable_dynamic_vae_sharding = (
         getattr(self.config, "enable_dynamic_vae_sharding", True) if hasattr(self, "config") else True
     )
-
     if enable_dynamic_vae_sharding and batch_size > 2:
       max_logging.log(
-          f"[Tuning] Skipping VAE replication and disabling slicing to prevent HBM OOM for batch_size {batch_size} > 2"
+          f"[Tuning] Disabling VAE slicing and applying dynamic batch sharding to prevent HBM OOM for batch_size {batch_size} > 2"
       )
       try:
         # Disable sequential slicing to avoid JAX concatenating 17GB arrays on the TPU
@@ -1847,14 +1846,21 @@ def convert_to_vel(lat, x0, sig):
         mesh = latents.sharding.mesh
         replicated_sharding = NamedSharding(mesh, P())
         latents = jax.lax.with_sharding_constraint(latents, replicated_sharding)
+      except Exception as e:  # pylint: disable=broad-exception-caught
+        max_logging.log(f"[Tuning] Failed to apply replicate VAE latents sharding: {e}")
+
+    if replicate_vae:
+      try:
+        mesh = latents.sharding.mesh
+        replicated_sharding = NamedSharding(mesh, P())
         # Replicate VAE weights
         graphdef, state = nnx.split(self.vae)
         state = jax.tree_util.tree_map(
             lambda x: jax.lax.with_sharding_constraint(x, replicated_sharding) if isinstance(x, jax.Array) else x, state
         )
         self.vae = nnx.merge(graphdef, state)
       except Exception as e:  # pylint: disable=broad-exception-caught
-        max_logging.log(f"[Tuning] Failed to apply sharding constraint: {e}")
+        max_logging.log(f"[Tuning] Failed to replicate VAE weights: {e}")
 
     latent_processing_time += time.perf_counter() - t0_latent_processing
     timings["Latent Processing"] = latent_processing_time