Fix multihost JAX Array fetching error in LTX2 pipelines by adding process_allgather before NumPy conversion

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1894,6 +1894,7 @@ def convert_to_vel(lat, x0, sig):
 
     # VAE outputs (B, T, H, W, C), but video processor expects (B, C, T, H, W)
     t0_video_post = time.perf_counter()
+    video = jax.experimental.multihost_utils.process_allgather(video, tiled=True)
     video_np = np.array(video).transpose(0, 4, 1, 2, 3)
     video = self.video_processor.postprocess_video(torch.from_numpy(video_np), output_type=output_type)
     video_post_time = time.perf_counter() - t0_video_post
@@ -1917,6 +1918,7 @@ def convert_to_vel(lat, x0, sig):
       audio = self._jitted_vocoder(self.vocoder, generated_mel_spectrograms)
 
     # Convert audio to numpy
+    audio = jax.experimental.multihost_utils.process_allgather(audio, tiled=True)
     audio = np.array(audio)
     vocoder_time = time.perf_counter() - t0_vocoder
     max_logging.log(f"Vocoder & Audio numpy time: {vocoder_time:.2f}s")

src/maxdiffusion/pipelines/ltx2/pipeline_ltx2_latent_upsample.py

@@ -250,6 +250,7 @@ def __call__(
     if video.dtype == jnp.bfloat16:
       video = video.astype(jnp.float32)
 
+    video = jax.experimental.multihost_utils.process_allgather(video, tiled=True)
     video = np.transpose(np.array(video), (0, 4, 1, 2, 3))
     video = self.video_processor.postprocess_video(video, output_type=output_type)