Add process_allgather to VAE decode outputs for multi-host execution

Author: amepas

src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py

@@ -298,6 +298,9 @@ def __call__(
     # VAE output is in decoded_out.sample
     images_rgb = decoded_out.sample
     images_rgb.block_until_ready()
+    if jax.process_count() > 1:
+      from jax.experimental import multihost_utils
+      images_rgb = multihost_utils.process_allgather(images_rgb, tiled=True)
 
     trace["vae_decode"] = time.perf_counter() - t0
     max_logging.log(f" -> [TIMING] VAE Decoding: {trace['vae_decode']:.4f} seconds ⏱️")