fix(wan): prevent duplicate batch concatenation on replicated VAE output

Author: jitendra-jalwaniya

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -940,7 +940,10 @@ def _decode_latents_to_video(self, latents: jax.Array, trace: Optional[dict] = N
       trace["vae_decode_tpu"] = time.perf_counter() - t_vae_tpu_start
 
     if hasattr(video, "addressable_shards") and len(video.addressable_shards) > 0:
-      video = np.concatenate([np.asarray(shard.data) for shard in video.addressable_shards], axis=0)
+      if video.addressable_shards[0].data.shape[0] < video.shape[0]:
+        video = np.concatenate([np.asarray(shard.data) for shard in video.addressable_shards], axis=0)
+      else:
+        video = np.asarray(video.addressable_shards[0].data)
     else:
       video = np.asarray(video)
     return video